Merge main into feat/runpod-serverless-workers

Resolve nodes/collector.py: keep the serverless manager accessor and
main's compressed result transfer. Update test stubs for the merged
code: fetch_worker_has_node (#99) in the serverless integration test,
and the real aiohttp.web in the transfer integration test.

.github/FUNDING.yml

@@ -1,4 +1,3 @@
 # These are supported funding model platforms
 
 github: robertvoy
-buy_me_a_coffee: robertvoy

.github/workflows/publish_action.yml

@@ -19,6 +19,6 @@ jobs:
       - name: Check out code
         uses: actions/checkout@v4
       - name: Publish Custom Node
-        uses: Comfy-Org/publish-node-action@v1
+        uses: Comfy-Org/publish-node-action@main
         with:
           personal_access_token: ${{ secrets.REGISTRY_ACCESS_TOKEN }}

.gitignore

@@ -7,8 +7,3 @@ __pycache__/
 node_modules/
 npm-debug.log*
 AGENTS.md
-.claude/
-
-# Desloppify artifacts
-.desloppify/
-scorecard.png

README.md

@@ -15,7 +15,7 @@
 
 ---
 
-## Key Features
+## Key Features
 
 #### Parallel Workflow Processing
 - Run your workflow on multiple GPUs simultaneously with varied seeds, collect results on the master
@@ -27,32 +27,23 @@
 - Intelligent distribution
 - Handles single images and videos
 
-#### Ease of Use
-- Auto-setup local workers; easily add remote/cloud ones
-- Convert any workflow to distributed with 2 nodes
-- JSON configuration with UI controls
-
----
-
-## Current Architecture
-
-- Workflow-level load balancing is controlled by **Distributed Collector** via the `load_balance` toggle.
-- There is **no Distributed Queue node** anymore.
-- With `load_balance=true`, orchestration selects one least-busy execution participant:
-  - If master participation is enabled, master is included as a candidate.
-  - If master is in orchestrator-only mode, only workers are considered.
-
----
+#### Ease of Use
+- Auto-setup local workers; easily add remote/cloud ones
+- Convert any workflow to distributed with 2 nodes
+- JSON configuration with UI controls
+
+---
 
 ## Worker Types
 
 <img width="200" align="right" alt="ComfyUI_temp_khvcc_00034_@0 25x" src="https://github.com/user-attachments/assets/651e4912-7c23-4e32-bd88-250f5175e129" />
 
-ComfyUI Distributed supports three types of workers:
+ComfyUI Distributed supports four types of workers:
 
 - **Local Workers** - Additional GPUs on the same machine (auto-configured on first launch)
 - **Remote Workers** - GPUs on other computers in your network
 - **Cloud Workers** - GPUs hosted on a cloud service like Runpod, accessible via secure tunnels
+- **Serverless (Runpod)** - Queue-based Runpod endpoints for image Collector jobs; no Distributed install or inbound tunnel on workers. [Setup, endpoint contract, and preview limitations](/docs/serverless-workers.md).
 
 > For detailed setup instructions, see the [setup guide](/docs/worker-setup-guides.md)
 
@@ -61,7 +52,6 @@ ComfyUI Distributed supports three types of workers:
 ## Requirements
 
 - ComfyUI
-> Note: Desktop app not currently supported
 - Multiple NVIDIA GPUs
 > No additional GPUs? Use [Cloud Workers](https://github.com/robertvoy/ComfyUI-Distributed/blob/main/docs/worker-setup-guides.md#cloud-workers)
 - That's it
@@ -92,19 +82,19 @@ Join Runpod with [this link](https://get.runpod.io/0bw29uf3ug0p) and unlock a sp
 
 ## Workflow Examples
 
-### Basic Parallel Generation
-Generate multiple images in the time it takes to generate one. Each worker uses a different seed.
+### Basic Parallel Generation
+Generate multiple images in the time it takes to generate one. Each worker uses a different seed.
 
 ![Clipboard Image (6)](https://github.com/user-attachments/assets/9598c94c-d9b4-4ccf-ab16-a21398220aeb)
 
 > [Download workflow](/workflows/distributed-txt2img.json)
 
 1. Open your ComfyUI workflow
-2. Add **Distributed Seed** → connect to sampler's seed
-3. Add **Distributed Collector** → after VAE Decode
-4. Optional: enable `load_balance` on Distributed Collector to run on one least-busy participant
-5. Enable workers in the UI
-6. Run the workflow!
+2. Add **Distributed Seed** → connect to sampler's seed
+3. Add **Distributed Collector** → after VAE Decode
+4. Optional: enable `load_balance` on Distributed Collector to run on one least-busy participant
+5. Enable workers in the UI
+6. Run the workflow!
 
 ### Parallel WAN Generation
 Generate multiple videos in the time it takes to generate one. Each worker uses a different seed.
@@ -122,12 +112,12 @@ Generate multiple videos in the time it takes to generate one. Each worker uses
 7. Enable workers in the UI
 8. Run the workflow!
 
-### Distributed Image Upscaling
-Accelerate Ultimate SD Upscaler by distributing tiles across multiple workers, with speed scaling as you add more GPUs.
+### Distributed Image Upscaling
+Accelerate Ultimate SD Upscaler by distributing tiles across multiple workers, with speed scaling as you add more GPUs.
 
 ![Clipboard Image (3)](https://github.com/user-attachments/assets/ffb57a0d-7b75-4497-96d2-875d60865a1a)
 
-> [Download workflow](/workflows/distributed-upscale.json)
+> [Download workflow](/workflows/distributed-upscale.json)
 
 1. Load your image
 2. Upscale with ESRGAN or similar
@@ -153,40 +143,40 @@ Accelerate Ultimate SD Upscaler by distributing video tiles across multiple work
 
 ---
 
-## Developer API
+## Developer API
 
 Control your distributed cluster programmatically without opening the browser.
 
-* **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)
+* **Endpoint:** `POST /distributed/queue`
+* **Functionality:** Accepts a ComfyUI API-format prompt, dispatches it to the requested reachable workers, and returns the master `prompt_id`.
+* **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).
-
----
-
-## Distributed Value
-
-Use **Distributed Value** when you want per-worker overrides (for example, different prompts/models/settings per worker).
-
-- Output type adapts to the connected input where possible (`STRING`, `INT`, `FLOAT`, `COMBO`).
-- The node shows only currently enabled workers.
-- If worker enablement changes, worker fields update automatically.
-- When disconnected, it resets to default string mode and clears per-worker overrides.
-- On execution, master uses `default_value`; workers use their mapped override with typed coercion fallback to default.
-
----
-
-## Nodes
-
-| Node | Description |
-|------|-------------|
-| **Distributed Seed** | Generates unique seeds for each worker |
-| **Distributed Collector** | Collects results (image/video frames and optionally audio) from workers back to the master; `load_balance` can route the run to one least-busy participant |
-| **Distributed Value** | Outputs per-worker override values with fallback to default |
-| **Ultimate SD Upscale Distributed** | Distributes upscale tiles across workers |
-| **Image Batch Divider** | Splits image batches for multi-GPU output |
-| **Audio Batch Divider** | Splits audio batches for multi-GPU output |
+> **⚠️ 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).
+
+---
+
+## Distributed Value
+
+Use **Distributed Value** when you want per-worker overrides (for example, different prompts/models/settings per worker).
+
+- Output type adapts to the connected input where possible (`STRING`, `INT`, `FLOAT`, `COMBO`).
+- The node shows only currently enabled workers.
+- If worker enablement changes, worker fields update automatically.
+- When disconnected, it resets to default string mode and clears per-worker overrides.
+- On execution, master uses `default_value`; workers use their mapped override with typed coercion fallback to default.
+
+---
+
+## Nodes
+
+| Node | Description |
+|------|-------------|
+| **Distributed Seed** | Generates unique seeds for each worker |
+| **Distributed Collector** | Collects results (image/video frames and optionally audio) from workers back to the master; `load_balance` can route the run to one least-busy participant |
+| **Distributed Value** | Outputs per-worker override values with fallback to default |
+| **Ultimate SD Upscale Distributed** | Distributes upscale tiles across workers |
+| **Image Batch Divider** | Splits image batches for multi-GPU output |
+| **Audio Segment Divider** | Splits an audio waveform into up to ten sequential time segments |
 | **Distributed Model Name** | Passes model paths to workers, enabling workflows to use models not present on the master in orchestrator-only mode |
 | **Distributed Empty Image** | Produces an empty IMAGE batch used when the master delegates all work |
 
@@ -206,7 +196,7 @@ No, it does not speed up the generation of a single image or video. Instead, it
 
 <details>
 <summary>Does it work with the ComfyUI desktop app?</summary>
-Currently, it is not compatible with the ComfyUI desktop app.
+Yes, it does now.
 </details>
 
 <details>
@@ -231,6 +221,26 @@ Yes, it is compatible with Docker setups, but you will need to configure your Do
 
 ---
 
+## Community projects
+
+Community-maintained add-ons and forks that extend ComfyUI-Distributed.
+
+### [ComfyUI-H3MC-Distributed](https://github.com/ymiyad/ComfyUI-H3MC-Distributed) by ymiyad
+
+Connects H3 Motion Context with ComfyUI-Distributed to generate multiple candidate takes of the next video clip across your workers. Compare the results, pick your preferred take, and use it as context for the next clip.
+
+Each saved candidate video includes its continuation latent, so you don't need a shared folder or separate latent files. This add-on speeds up iteration through parallel candidates; it does not split a single clip across GPUs.
+
+### [ComfyUI-Deadline-Distributed](https://github.com/doubletwisted/ComfyUI-Deadline-Distributed) by doubletwisted
+
+A fork of ComfyUI-Distributed with Thinkbox/AWS Deadline 10 integration for render farms. Claim Deadline workers from the Distributed panel, register them automatically, and collect their results back into your master workflow.
+
+Requires Deadline 10 and the author's ComfyUI-Deadline-Plugin.
+
+See each repository for installation instructions, requirements, and support.
+
+---
+
 ## Disclaimer
 
 This software is provided "as is" without any warranties, express or implied, including merchantability, fitness for a particular purpose, or non-infringement. The developers and copyright holders are not liable for any claims, damages, or liabilities arising from the use, modification, or distribution of the software. Users are solely responsible for ensuring compliance with applicable laws and regulations and for securing their networks against unauthorized access, hacking, data breaches, or loss. The developers assume no liability for any damages or incidents resulting from misuse, improper configuration, or external threats.
@@ -249,3 +259,4 @@ Buy me a coffee at: https://buymeacoffee.com/robertvoy
 
 
 
+

__init__.py

@@ -1,11 +1,22 @@
-"""ComfyUI-Distributed package entrypoint."""
-from __future__ import annotations
+"""ComfyUI-Distributed's native V3 extension entrypoint."""
+from comfy_api.v0_0_2 import ComfyExtension, io
 
-from .bootstrap.entrypoint import build_node_mappings, initialize_runtime
+from .nodes.v3 import NODES
+from .runtime.bootstrap import initialize
 
-WEB_DIRECTORY = "./web"
+WEB_DIRECTORY = './web'
 
-NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS = build_node_mappings()
-initialize_runtime()
 
-__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS", "WEB_DIRECTORY"]
+class DistributedExtension(ComfyExtension):
+    async def on_load(self) -> None:
+        initialize()
+
+    async def get_node_list(self) -> list[type[io.ComfyNode]]:
+        return list(NODES)
+
+
+async def comfy_entrypoint() -> DistributedExtension:
+    return DistributedExtension()
+
+
+__all__ = ['comfy_entrypoint', 'WEB_DIRECTORY']

api/__init__.py

@@ -1,21 +1,6 @@
-"""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"]
+from . import config_routes  # noqa: F401
+from . import serverless_routes  # noqa: F401
+from . import tunnel_routes  # noqa: F401
+from . import worker_routes  # noqa: F401
+from . import job_routes  # noqa: F401
+from . import usdu_routes  # noqa: F401

api/config_routes.py

@@ -1,15 +1,39 @@
-from __future__ import annotations
+import json
+from contextlib import asynccontextmanager
+from typing import Any
 
 from aiohttp import web
 import server
 
-from ..utils.config import config_transaction, load_config
-from ..utils.logging import debug_log
+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.network import handle_api_error, normalize_host
-from .endpoint_policy import run_authorized_endpoint
+from ..utils.secret_fields import is_secret_field_name
+from ..utils.serverless_config import (
+    ServerlessSecretChanges,
+    public_config,
+    validate_serverless_worker,
+)
 
 
-def _positive_int(value: int) -> bool:
+def _positive_int(value):
     return value > 0
 
 
@@ -69,103 +93,237 @@ 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"]
+_SECRET_FIELDS = {"api_key", "clear_api_key"}
+_SERVERLESS_LOCAL_FIELDS = {"host", "port", "listen_port", "cuda_device", "extra_args"}
+_SERVERLESS_ONLY_FIELDS = {
+    "provider", "endpoint_id", "jobs_per_run", "concurrency", "timeout_seconds", "api_key_configured",
+}
+_SERVERLESS_CONFIG_FIELDS = {
+    "worker_id", "type", "provider", "endpoint_id", "jobs_per_run", "concurrency",
+    "timeout_seconds", "enabled", "name", "api_key", "clear_api_key",
+}
 
-    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}'")
+class _ActivePoolCheckError(RuntimeError):
+    pass
+
+
+def _error_response(message, status):
+    return web.json_response({"status": "error", "message": message}, status=status)
+
+
+def _contains_secret_fields(value, *, allow_root_serverless_secrets=False, depth=0):
+    if isinstance(value, dict):
+        for key, item in value.items():
+            field_name = key.casefold() if isinstance(key, str) else ""
+            if is_secret_field_name(field_name):
+                if not (allow_root_serverless_secrets and depth == 0 and key in _SECRET_FIELDS):
+                    return True
+            if _contains_secret_fields(
+                item, allow_root_serverless_secrets=allow_root_serverless_secrets, depth=depth + 1
+            ):
+                return True
+    elif isinstance(value, (list, tuple)):
+        return any(
+            _contains_secret_fields(
+                item, allow_root_serverless_secrets=allow_root_serverless_secrets, depth=depth + 1
+            )
+            for item in value
+        )
+    return False
+
+
+def _serverless_secret_store():
+    from .orchestration.serverless_runtime import get_secret_store
+    return get_secret_store()
+
+
+def _serverless_pool_is_active(worker_id):
+    from .orchestration.serverless_runtime import get_manager
+    return get_manager().pool_is_active(worker_id)
+
+
+def _ensure_pool_idle(worker_id):
+    try:
+        active = _serverless_pool_is_active(worker_id)
+    except Exception:
+        raise _ActivePoolCheckError from None
+    if active:
+        raise _ActivePoolCheckError
+
+
+def _worker_by_id(workers, worker_id):
+    for worker in workers:
+        if isinstance(worker, dict) and worker.get("id") == worker_id:
+            return worker
+    return None
+
+
+def _worker_type(worker):
+    return worker.get("type", "local") if isinstance(worker, dict) else None
+
+
+def _strip_sensitive_fields(value: Any) -> Any:
+    if isinstance(value, dict):
+        return {
+            key: _strip_sensitive_fields(item)
+            for key, item in value.items()
+            if not is_secret_field_name(key)
+        }
+    if isinstance(value, list):
+        return [_strip_sensitive_fields(item) for item in value]
+    return value
+
+
+def _clean_serverless_worker(worker):
+    clean = _strip_sensitive_fields(dict(worker))
+    for field in _SERVERLESS_LOCAL_FIELDS | {"api_key_configured"}:
+        clean.pop(field, None)
+    return clean
+
+
+def _validate_bulk_workers(workers):
+    validated = []
+    for worker in workers:
+        if isinstance(worker, dict) and worker.get("type") == "serverless":
+            clean = _clean_serverless_worker(worker)
+            validated.append(validate_serverless_worker(clean))
         else:
-            seen_ids.add(worker_id)
+            validated.append(worker)
+    return validated
 
-        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
+def _clean_bulk_type_conversions(old_workers, new_workers):
+    old_by_id = {str(w.get("id")): w for w in old_workers if isinstance(w, dict) and w.get("id") is not None}
+    for index, worker in enumerate(new_workers):
+        if not isinstance(worker, dict) or worker.get("id") is None:
+            continue
+        old = old_by_id.get(str(worker["id"]))
+        if old is None or old.get("type") != "serverless" or worker.get("type") == "serverless":
+            continue
+        converted = dict(worker)
+        for field in _SERVERLESS_ONLY_FIELDS:
+            converted.pop(field, None)
+        new_workers[index] = converted
+    return new_workers
+
+
+def _bulk_serverless_changes(old_workers, new_workers):
+    old_by_id = {str(w.get("id")): w for w in old_workers if isinstance(w, dict) and w.get("id") is not None}
+    new_by_id = {str(w.get("id")): w for w in new_workers if isinstance(w, dict) and w.get("id") is not None}
+    secret_ids = set()
+    for worker_id, old in old_by_id.items():
+        new = new_by_id.get(worker_id)
+        if old.get("type") == "serverless":
+            if new is None or new.get("type") != "serverless":
+                _ensure_pool_idle(worker_id)
+                secret_ids.add(worker_id)
+            elif old.get("endpoint_id") != new.get("endpoint_id"):
+                _ensure_pool_idle(worker_id)
+        elif new is not None and new.get("type") == "serverless":
+            secret_ids.add(worker_id)
+    return secret_ids
 
 
 @server.PromptServer.instance.routes.get("/distributed/config")
-async def get_config_endpoint(request: web.Request) -> web.StreamResponse:
-    async def _operation() -> web.StreamResponse:
-        return web.json_response(load_config())
-
-    return await run_authorized_endpoint(request, _operation)
+async def get_config_endpoint(request):
+    try:
+        config = load_config()
+        return web.json_response(public_config(config, _serverless_secret_store()))
+    except Exception:
+        return _error_response("Unable to read config", 500)
 
 
 @server.PromptServer.instance.routes.post("/distributed/config")
-async def update_config_endpoint(request: web.Request) -> web.StreamResponse:
-    """Bulk config update with schema validation."""
-    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)
+async def update_config_endpoint(request):
+    """Bulk config update with schema validation and credential redaction."""
+    try:
+        data = await request.json()
+    except Exception:
+        return _error_response("Invalid JSON payload", 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 _error_response("Config payload must be an object", 400)
+    if _contains_secret_fields(data):
+        return _error_response("API keys cannot be set through bulk config", 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}: failed validation")
+            continue
 
-            if key == "workers":
-                errors.extend(_validate_workers_payload(value))
-                if errors:
-                    continue
+        if key == "workers":
+            try:
+                value = _validate_bulk_workers(value)
+            except (TypeError, ValueError):
+                return _error_response("Invalid serverless worker configuration", 400)
 
-            if key in _SETTINGS_FIELDS:
-                validated_settings[key] = value
-            else:
-                validated_root[key] = value
+        if key in _SETTINGS_FIELDS:
+            validated_settings[key] = value
+        else:
+            validated_root[key] = value
 
-        if errors:
-            return await handle_api_error(request, errors, 400)
+    if errors:
+        return web.json_response({
+            "status": "error",
+            "error": errors,
+            "message": "; ".join(errors),
+        }, status=400)
 
-        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})
+    store = None
+    response_config = None
+    try:
+        with ServerlessSecretChanges() as secret_changes:
+            async with config_transaction() as config:
+                if "workers" in validated_root:
+                    old_workers = config.get("workers", [])
+                    validated_root["workers"] = _clean_bulk_type_conversions(
+                        old_workers, validated_root["workers"]
+                    )
+                    changed_ids = _bulk_serverless_changes(old_workers, validated_root["workers"])
+                    for worker in validated_root["workers"]:
+                        if isinstance(worker, dict) and worker.get("type") == "serverless" and worker.get("enabled", False):
+                            store = store or _serverless_secret_store()
+                            if worker["id"] in changed_ids or not store.has(worker["id"]):
+                                return _error_response("Save an API key before enabling this worker", 400)
+                    if changed_ids:
+                        store = store or _serverless_secret_store()
+                        for worker_id in sorted(changed_ids):
+                            prior_secret = secret_changes.capture(store, worker_id)
+                            if prior_secret is not None:
+                                secret_changes.delete(store, worker_id)
 
-    return await run_authorized_endpoint(request, _operation)
+                settings = config.setdefault("settings", {})
+                settings.update(validated_settings)
+                for key, value in validated_root.items():
+                    config[key] = value
+                response_config = public_config(config, store or _serverless_secret_store())
+            response = web.json_response({"status": "success", "config": response_config})
+            secret_changes.commit()
+            return response
+    except _ActivePoolCheckError:
+        return _error_response("Finish or reconcile unresolved serverless jobs before changing this worker", 409)
+    except Exception:
+        return _error_response("Unable to update config", 500)
 
 
 @server.PromptServer.instance.routes.get("/distributed/queue_status/{job_id}")
-async def queue_status_endpoint(request: web.Request) -> web.StreamResponse:
+async def queue_status_endpoint(request):
     """Check if a job queue is initialized."""
-    async def _operation() -> web.StreamResponse:
+    try:
         job_id = request.match_info['job_id']
         
         # Import to ensure initialization
@@ -177,89 +335,212 @@ 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})
-
-    return await run_authorized_endpoint(request, _operation, unexpected_status=500)
+    except Exception as e:
+        return await handle_api_error(request, e, 500)
 
 @server.PromptServer.instance.routes.post("/distributed/config/update_worker")
-async def update_worker_endpoint(request: web.Request) -> web.StreamResponse:
-    async def _operation() -> web.StreamResponse:
-        data = await request.json()
+async def update_worker_endpoint(request):
+    return await _run_with_serverless_secret_changes(_update_worker_endpoint, request)
+
+
+async def _run_with_serverless_secret_changes(endpoint, request):
+    with ServerlessSecretChanges() as secret_changes:
+        response = await endpoint(request, secret_changes)
+        if response.status < 400:
+            secret_changes.commit()
+        return response
+
+
+async def _update_worker_endpoint(request, secret_changes):
+    store = None
+    worker_id = None
+    try:
+        try:
+            data = await request.json()
+        except Exception:
+            return _error_response("Invalid JSON payload", 400)
+        if not isinstance(data, dict):
+            return _error_response("Worker payload must be an object", 400)
+
         worker_id = data.get("worker_id")
-        
         if worker_id is None:
-            return await handle_api_error(request, "Missing worker_id", 400)
+            return _error_response("Missing worker_id", 400)
 
+        if _contains_secret_fields(data, allow_root_serverless_secrets=True):
+            return _error_response("Invalid worker credentials", 400)
+
+        result = None
         async with config_transaction() as config:
-            worker_found = False
             workers = config.setdefault("workers", [])
+            existing = _worker_by_id(workers, worker_id)
+            current_type = _worker_type(existing)
+            target_type = data.get("type", current_type or "local")
 
-            for worker in workers:
-                if worker["id"] == worker_id:
-                    _apply_field_patch(worker, data, _WORKER_FIELDS)
-                    worker_found = True
-                    break
+            if target_type == "serverless":
+                unknown = set(data) - _SERVERLESS_CONFIG_FIELDS - _SERVERLESS_LOCAL_FIELDS
+                if unknown:
+                    return _error_response("Unknown serverless worker field", 400)
+                if not isinstance(worker_id, str) or not worker_id.strip():
+                    return _error_response("worker_id must be a non-empty string", 400)
+                if "api_key" in data and not isinstance(data["api_key"], str):
+                    return _error_response("Invalid worker credentials", 400)
+                if "clear_api_key" in data and type(data["clear_api_key"]) is not bool:
+                    return _error_response("Invalid worker credentials", 400)
+                supplied_key = data.get("api_key", "")
+                clear_key = data.get("clear_api_key", False)
+                if supplied_key and clear_key:
+                    return _error_response("Choose either api_key or clear_api_key", 400)
+                if "enabled" in data and type(data["enabled"]) is not bool:
+                    return _error_response("enabled must be a boolean", 400)
+                if "name" in data and not isinstance(data["name"], str):
+                    return _error_response("name must be a string", 400)
 
-            if not worker_found:
-                # If worker not found and all required fields are provided, create new worker
-                if all(key in data for key in ["name", "port", "cuda_device"]):
-                    new_worker = {
-                        "id": worker_id,
-                        "name": data["name"],
-                        "host": normalize_host(data.get("host", "localhost")),
-                        "port": data["port"],
-                        "cuda_device": data["cuda_device"],
-                        "enabled": data.get("enabled", False),
-                        "extra_args": data.get("extra_args", ""),
-                        "type": data.get("type", "local")
-                    }
-                    workers.append(new_worker)
+                candidate = dict(existing or {})
+                candidate["id"] = worker_id
+                candidate["type"] = "serverless"
+                candidate.setdefault("provider", "runpod")
+                candidate.setdefault("enabled", False)
+                # Legacy "concurrency" is accepted (see _SERVERLESS_CONFIG_FIELDS) but ignored.
+                for field in ("provider", "endpoint_id", "jobs_per_run", "timeout_seconds", "enabled", "name"):
+                    if field in data:
+                        candidate[field] = data[field]
+                candidate = _clean_serverless_worker(candidate)
+                validated = validate_serverless_worker(candidate)
+                if "enabled" in validated and type(validated["enabled"]) is not bool:
+                    return _error_response("enabled must be a boolean", 400)
+                if "name" in validated and not isinstance(validated["name"], str):
+                    return _error_response("name must be a string", 400)
+
+                store = _serverless_secret_store()
+                prior_secret = secret_changes.capture(store, worker_id)
+                converting_to_serverless = current_type != "serverless"
+                next_secret = None if converting_to_serverless else prior_secret
+                if supplied_key:
+                    next_secret = supplied_key
+                elif clear_key:
+                    next_secret = None
+
+                endpoint_changed = (
+                    current_type == "serverless"
+                    and existing.get("endpoint_id") != validated.get("endpoint_id")
+                )
+                if current_type == "serverless" and (
+                    endpoint_changed or next_secret != prior_secret
+                ):
+                    _ensure_pool_idle(worker_id)
+
+                if validated.get("enabled", False) and next_secret is None:
+                    return _error_response("Save an API key before enabling this worker; turn the worker off before clearing its key", 400)
+
+                if next_secret != prior_secret:
+                    if next_secret is None:
+                        secret_changes.delete(store, worker_id)
+                    else:
+                        secret_changes.set(store, worker_id, next_secret)
+
+                if existing is None:
+                    workers.append(validated)
                 else:
-                    return await handle_api_error(
-                        request,
-                        f"Worker {worker_id} not found and missing required fields for creation",
-                        404,
-                    )
+                    existing.clear()
+                    existing.update(validated)
+                result = {"status": "success", "api_key_configured": next_secret is not None}
 
-            return web.json_response({"status": "success"})
+            elif existing is not None:
+                cleaned_existing = _strip_sensitive_fields(existing)
+                existing.clear()
+                existing.update(cleaned_existing)
+                if current_type == "serverless":
+                    if _contains_secret_fields(data):
+                        return _error_response("Invalid worker credentials", 400)
+                    _ensure_pool_idle(worker_id)
+                    store = _serverless_secret_store()
+                    prior_secret = secret_changes.capture(store, worker_id)
+                    if prior_secret is not None:
+                        secret_changes.delete(store, worker_id)
+                    for field in _SERVERLESS_ONLY_FIELDS:
+                        existing.pop(field, None)
+                elif _contains_secret_fields(data):
+                    return _error_response("API keys are only accepted for serverless workers", 400)
+                _apply_field_patch(existing, data, _WORKER_FIELDS)
+                result = {"status": "success"}
+            elif all(key in data for key in ["name", "port", "cuda_device"]):
+                if _contains_secret_fields(data):
+                    return _error_response("API keys are only accepted for serverless workers", 400)
+                new_worker = {
+                    "id": worker_id,
+                    "name": data["name"],
+                    "host": normalize_host(data.get("host", "localhost")),
+                    "port": data["port"],
+                    "cuda_device": data["cuda_device"],
+                    "enabled": data.get("enabled", False),
+                    "extra_args": data.get("extra_args", ""),
+                    "type": data.get("type", "local")
+                }
+                workers.append(new_worker)
+                result = {"status": "success"}
+            else:
+                return _error_response(
+                    "Worker not found and missing required fields for creation", 404
+                )
+
+        return web.json_response(result)
+    except _ActivePoolCheckError:
+        return _error_response("Finish or reconcile unresolved serverless jobs before changing this worker", 409)
+    except ValueError:
+        return _error_response("Invalid worker configuration", 400)
+    except Exception:
+        return _error_response("Unable to update worker configuration", 500)
 
-    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:
-    async def _operation() -> web.StreamResponse:
-        data = await request.json()
+async def delete_worker_endpoint(request):
+    return await _run_with_serverless_secret_changes(_delete_worker_endpoint, request)
+
+
+async def _delete_worker_endpoint(request, secret_changes):
+    store = None
+    worker_id = None
+    try:
+        try:
+            data = await request.json()
+        except Exception:
+            return _error_response("Invalid JSON payload", 400)
+        if not isinstance(data, dict):
+            return _error_response("Worker payload must be an object", 400)
         worker_id = data.get("worker_id")
-        
         if worker_id is None:
-            return await handle_api_error(request, "Missing worker_id", 400)
-            
+            return _error_response("Missing worker_id", 400)
+
+        result = None
         async with config_transaction() as config:
             workers = config.get("workers", [])
-
-            # Find and remove the worker
-            worker_index = -1
-            for i, worker in enumerate(workers):
-                if worker["id"] == worker_id:
-                    worker_index = i
-                    break
-
-            if worker_index == -1:
-                return await handle_api_error(request, f"Worker {worker_id} not found", 404)
-
-            # Remove the worker
-            removed_worker = workers.pop(worker_index)
-
-            return web.json_response({
+            removed = _worker_by_id(workers, worker_id)
+            if removed is None:
+                return _error_response("Worker not found", 404)
+            if _worker_type(removed) == "serverless":
+                _ensure_pool_idle(worker_id)
+                store = _serverless_secret_store()
+                secret_worker_id = str(worker_id)
+                prior_secret = secret_changes.capture(store, secret_worker_id)
+                if prior_secret is not None:
+                    secret_changes.delete(store, secret_worker_id)
+            workers.remove(removed)
+            result = {
                 "status": "success",
-                "message": f"Worker {removed_worker.get('name', worker_id)} deleted"
-            })
+                "message": "Worker deleted" if _worker_type(removed) == "serverless"
+                else f"Worker {removed.get('name', worker_id)} deleted",
+            }
 
-    return await run_authorized_endpoint(request, _operation, unexpected_status=500)
+        return web.json_response(result)
+    except _ActivePoolCheckError:
+        return _error_response("Finish or reconcile unresolved serverless jobs before changing this worker", 409)
+    except Exception:
+        return _error_response("Unable to delete worker", 500)
 
 @server.PromptServer.instance.routes.post("/distributed/config/update_setting")
-async def update_setting_endpoint(request: web.Request) -> web.StreamResponse:
+async def update_setting_endpoint(request):
     """Updates a specific key in the settings object."""
-    async def _operation() -> web.StreamResponse:
+    try:
         data = await request.json()
         key = data.get("key")
         value = data.get("value")
@@ -276,13 +557,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."})
-
-    return await run_authorized_endpoint(request, _operation, unexpected_status=500)
+    except Exception as e:
+        return await handle_api_error(request, e, 400)
 
 @server.PromptServer.instance.routes.post("/distributed/config/update_master")
-async def update_master_endpoint(request: web.Request) -> web.StreamResponse:
+async def update_master_endpoint(request):
     """Updates master configuration."""
-    async def _operation() -> web.StreamResponse:
+    try:
         data = await request.json()
         
         async with config_transaction() as config:
@@ -291,5 +572,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."})
-
-    return await run_authorized_endpoint(request, _operation, unexpected_status=500)
+    except Exception as e:
+        return await handle_api_error(request, e, 400)

api/endpoint_policy.py

@@ -1,25 +0,0 @@
-"""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

@@ -1,5 +1,3 @@
-from __future__ import annotations
-
 import json
 import asyncio
 import io
@@ -7,7 +5,6 @@ import os
 import base64
 import binascii
 import time
-from typing import Any
 
 from aiohttp import web
 import server
@@ -18,21 +15,31 @@ 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
-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 ..utils import transfer_codec
+try:
+    from .queue_orchestration import ensure_distributed_state, orchestrate_distributed_execution
+except ImportError:
+    from .queue_orchestration import orchestrate_distributed_execution
+
+    def ensure_distributed_state():
+        return None
 from .queue_request import parse_queue_request_payload
 
-# 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 }
+prompt_server = server.PromptServer.instance
+
+# Worker result envelopes accepted by POST /distributed/job_complete:
+#
+# application/json (legacy, one frame per request):
+#   { "job_id": str, "worker_id": str, "batch_idx": int, "image": <base64 PNG>, "is_last": bool,
+#     "audio"?: {...} }
+#
+# multipart/form-data (compressed, many frames per request):
+#   envelope: JSON { "job_id", "worker_id", "format", "start_index", "count", "width", "height",
+#                    "is_last", "audio"?: {...} }
+#   media:    one part per frame for still formats, or one video file for video formats
 
 
-def _runtime_state():
-    return ensure_distributed_runtime_state()
-
-
-def _decode_image_sync(image_path: str) -> dict[str, Any]:
+def _decode_image_sync(image_path):
     """Decode image/video file and compute hash in a threadpool worker."""
     import base64
     import hashlib
@@ -78,7 +85,7 @@ def _decode_image_sync(image_path: str) -> dict[str, Any]:
     }
 
 
-def _check_file_sync(filename: str, expected_hash: str) -> dict[str, Any]:
+def _check_file_sync(filename, expected_hash):
     """Check file presence and hash in a threadpool worker."""
     import hashlib
     import folder_paths
@@ -103,7 +110,7 @@ def _check_file_sync(filename: str, expected_hash: str) -> dict[str, Any]:
     }
 
 
-def _decode_canonical_png_tensor(image_payload: str) -> torch.Tensor:
+def _decode_canonical_png_tensor(image_payload):
     """Decode canonical base64 PNG payload into a contiguous IMAGE tensor."""
     if not isinstance(image_payload, str) or not image_payload.strip():
         raise ValueError("Field 'image' must be a non-empty base64 PNG string.")
@@ -134,7 +141,7 @@ def _decode_canonical_png_tensor(image_payload: str) -> torch.Tensor:
         raise ValueError(f"Failed to decode PNG image payload: {exc}") from exc
 
 
-def _decode_audio_payload(audio_payload: dict[str, Any]) -> dict[str, Any]:
+def _decode_audio_payload(audio_payload):
     """Decode canonical audio payload into an AUDIO dict."""
     from ..utils.audio_payload import decode_audio_payload
 
@@ -142,20 +149,17 @@ 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
+async def prepare_job_endpoint(request):
     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)
 
-        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()
+        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()
         
         debug_log(f"Prepared queue for job {multi_job_id}")
         return web.json_response({"status": "success"})
@@ -163,13 +167,9 @@ async def prepare_job_endpoint(request: web.Request) -> web.StreamResponse:
         return await handle_api_error(request, e)
 
 @server.PromptServer.instance.routes.post("/distributed/clear_memory")
-async def clear_memory_endpoint(request: web.Request) -> web.StreamResponse:
+async def clear_memory_endpoint(request):
     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)
@@ -186,16 +186,12 @@ async def clear_memory_endpoint(request: web.Request) -> web.StreamResponse:
         try:
             mm.unload_all_models()
         except AttributeError as e:
-            warning = f"Model unload warning: {e}"
-            warnings.append(warning)
-            debug_log(warning)
+            debug_log(f"Warning during model unload: {e}")
         
         try:
             mm.soft_empty_cache()
         except Exception as e:
-            warning = f"Cache clear warning: {e}"
-            warnings.append(warning)
-            debug_log(warning)
+            debug_log(f"Warning during cache clear: {e}")
         
         for _ in range(3):
             gc.collect()
@@ -204,16 +200,6 @@ 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:
@@ -222,16 +208,13 @@ async def clear_memory_endpoint(request: web.Request) -> web.StreamResponse:
         gc.collect()
         if torch.cuda.is_available():
             torch.cuda.empty_cache()
-        debug_log(f"VRAM clear failed: {e}")
-        return await handle_api_error(request, f"GPU memory clear failed: {e}", 500)
+        debug_log(f"Partial VRAM clear completed with warning: {e}")
+        return web.json_response({"status": "success", "message": "GPU memory cleared (with warnings)"})
 
 
 @server.PromptServer.instance.routes.post("/distributed/queue")
-async def distributed_queue_endpoint(request: web.Request) -> web.StreamResponse:
+async def distributed_queue_endpoint(request):
     """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:
@@ -243,7 +226,7 @@ async def distributed_queue_endpoint(request: web.Request) -> web.StreamResponse
         return await handle_api_error(request, exc, 400)
 
     try:
-        prompt_id, worker_count = await orchestrate_distributed_execution(
+        prompt_id, prompt_number, worker_count, node_errors = await orchestrate_distributed_execution(
             payload.prompt,
             payload.workflow_meta,
             payload.client_id,
@@ -253,17 +236,17 @@ async def distributed_queue_endpoint(request: web.Request) -> web.StreamResponse
         )
         return web.json_response({
             "prompt_id": prompt_id,
+            "number": prompt_number,
+            "node_errors": node_errors,
             "worker_count": worker_count,
+            "auto_prepare_supported": True,
         })
     except Exception as exc:
         return await handle_api_error(request, exc, 500)
 
 @server.PromptServer.instance.routes.post("/distributed/load_image")
-async def load_image_endpoint(request: web.Request) -> web.StreamResponse:
+async def load_image_endpoint(request):
     """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")
@@ -279,11 +262,8 @@ async def load_image_endpoint(request: web.Request) -> web.StreamResponse:
         return await handle_api_error(request, e, 500)
 
 @server.PromptServer.instance.routes.post("/distributed/check_file")
-async def check_file_endpoint(request: web.Request) -> web.StreamResponse:
+async def check_file_endpoint(request):
     """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")
@@ -299,11 +279,136 @@ async def check_file_endpoint(request: web.Request) -> web.StreamResponse:
         return await handle_api_error(request, e, 500)
 
 
+async def _enqueue_worker_items(multi_job_id, items):
+    """Put decoded items on the job queue, waiting briefly for the master to create it.
+
+    Returns the queue size after enqueueing, or None if the job never appeared.
+    """
+    deadline = time.monotonic() + float(JOB_INIT_GRACE_PERIOD)
+    while True:
+        async with prompt_server.distributed_jobs_lock:
+            pending = prompt_server.distributed_pending_jobs.get(multi_job_id)
+            if pending is not None:
+                for item in items:
+                    await pending.put(item)
+                return pending.qsize()
+        if time.monotonic() > deadline:
+            return None
+        await asyncio.sleep(0.05)
+
+
+def _decode_compressed_chunk_sync(envelope, parts):
+    """Decode a compressed chunk into per-frame IMAGE tensors (threadpool)."""
+    frames = transfer_codec.decode_chunk(envelope, parts)
+    return [
+        ensure_contiguous(torch.from_numpy(transfer_codec.uint8_frame_to_tensor_array(frame)))
+        for frame in frames
+    ]
+
+
+def _validate_compressed_envelope(envelope):
+    errors = []
+    if not isinstance(envelope, dict):
+        return ["envelope: expected JSON object"]
+    for key in ("job_id", "worker_id"):
+        value = envelope.get(key)
+        if not isinstance(value, str) or not value.strip():
+            errors.append(f"{key}: expected non-empty string")
+    start_index = envelope.get("start_index")
+    if not isinstance(start_index, int) or isinstance(start_index, bool) or start_index < 0:
+        errors.append("start_index: expected non-negative integer")
+    if not isinstance(envelope.get("is_last"), bool):
+        errors.append("is_last: expected boolean")
+    try:
+        if transfer_codec.normalize_format(envelope.get("format")) == transfer_codec.LEGACY_FORMAT:
+            errors.append("format: legacy_png must be sent as JSON")
+    except transfer_codec.TransferCodecError as exc:
+        errors.append(f"format: {exc}")
+    audio_payload = envelope.get("audio")
+    if audio_payload is not None and not isinstance(audio_payload, dict):
+        errors.append("audio: expected object when provided")
+    return errors
+
+
+async def _job_complete_compressed(request):
+    try:
+        form = await request.post()
+    except Exception as exc:
+        too_large = getattr(web, "HTTPRequestEntityTooLarge", None)
+        if too_large is not None and isinstance(exc, too_large):
+            return await handle_api_error(
+                request, f"Request body exceeds the master's upload limit (--max-upload-size): {exc}", 413)
+        return await handle_api_error(request, f"Invalid multipart payload: {exc}", 400)
+
+    raw_envelope = form.get("envelope")
+    if isinstance(raw_envelope, (bytes, bytearray)):
+        raw_envelope = raw_envelope.decode("utf-8")
+    elif hasattr(raw_envelope, "file"):
+        raw_envelope = raw_envelope.file.read().decode("utf-8")
+    try:
+        envelope = json.loads(raw_envelope) if isinstance(raw_envelope, str) else None
+    except ValueError as exc:
+        return await handle_api_error(request, f"envelope: invalid JSON ({exc})", 400)
+    if envelope is None:
+        return await handle_api_error(request, "envelope: missing field", 400)
+
+    errors = _validate_compressed_envelope(envelope)
+    if errors:
+        return await handle_api_error(request, errors, 400)
+
+    parts = []
+    for field in form.getall("media", []):
+        if hasattr(field, "file"):
+            parts.append(field.file.read())
+        elif isinstance(field, (bytes, bytearray)):
+            parts.append(bytes(field))
+        else:
+            return await handle_api_error(request, "media: expected binary file parts", 400)
+
+    loop = asyncio.get_running_loop()
+    try:
+        tensors = await loop.run_in_executor(None, _decode_compressed_chunk_sync, envelope, parts)
+        decoded_audio = _decode_audio_payload(envelope.get("audio")) if envelope.get("audio") is not None else None
+    except ValueError as exc:
+        return await handle_api_error(request, exc, 400)
+
+    multi_job_id = envelope["job_id"].strip()
+    worker_id = envelope["worker_id"].strip()
+    start_index = envelope["start_index"]
+    is_last = envelope["is_last"]
+    last_pos = len(tensors) - 1
+    items = [
+        {
+            "tensor": tensor,
+            "worker_id": worker_id,
+            "image_index": start_index + pos,
+            "is_last": bool(is_last and pos == last_pos),
+            "audio": decoded_audio if pos == last_pos else None,
+        }
+        for pos, tensor in enumerate(tensors)
+    ]
+
+    queue_size = await _enqueue_worker_items(multi_job_id, items)
+    if queue_size is None:
+        return await handle_api_error(request, "job not initialized", 404)
+
+    debug_log(
+        f"job_complete received {envelope['format']} chunk - job_id: {multi_job_id}, worker: {worker_id}, "
+        f"frames: {start_index}..{start_index + len(items) - 1}, "
+        f"bytes: {sum(len(p) for p in parts)}, is_last: {is_last}, queue_size: {queue_size}"
+    )
+    return web.json_response({"status": "success", "frames": len(items)})
+
+
 @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
+async def job_complete_endpoint(request):
+    content_type = str(getattr(request, "content_type", "") or "").lower()
+    if content_type.startswith("multipart/"):
+        try:
+            return await _job_complete_compressed(request)
+        except Exception as e:
+            return await handle_api_error(request, e)
+
     try:
         data = await request.json()
     except Exception as exc:
@@ -318,7 +423,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_raw = data.get("is_last")
+        is_last = data.get("is_last")
 
         errors = []
         if not isinstance(job_id, str) or not job_id.strip():
@@ -327,57 +432,41 @@ async def job_complete_endpoint(request: web.Request) -> web.StreamResponse:
             errors.append("worker_id: expected non-empty string")
         if not isinstance(batch_idx, int) or batch_idx < 0:
             errors.append("batch_idx: expected non-negative integer")
-        if not isinstance(image_payload, str) or not image_payload.strip():
+        image_provided = image_payload is not None
+        audio_provided = audio_payload is not None
+        if not image_provided and not audio_provided:
+            errors.append("expected at least one of image or audio")
+        if image_provided and (not isinstance(image_payload, str) or not image_payload.strip()):
             errors.append("image: expected non-empty base64 PNG string")
-        if audio_payload is not None and not isinstance(audio_payload, dict):
+        if audio_provided and not isinstance(audio_payload, dict):
             errors.append("audio: expected object when provided")
-
-        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 not isinstance(is_last, bool):
+            errors.append("is_last: expected boolean")
         if errors:
             return await handle_api_error(request, errors, 400)
 
-        tensor = _decode_canonical_png_tensor(image_payload)
-        decoded_audio = _decode_audio_payload(audio_payload) if audio_payload is not None else None
+        try:
+            tensor = _decode_canonical_png_tensor(image_payload) if image_provided else None
+            decoded_audio = _decode_audio_payload(audio_payload) if audio_provided else None
+        except ValueError as exc:
+            return await handle_api_error(request, exc, 400)
         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 runtime_state.distributed_jobs_lock:
-                pending = runtime_state.distributed_pending_jobs.get(multi_job_id)
-                if pending is not None:
-                    queue_item = {
-                        "tensor": tensor,
-                        "worker_id": worker_id,
-                        "image_index": int(batch_idx),
-                        "is_last": is_last,
-                    }
-                    if decoded_audio is not None:
-                        queue_item["audio"] = decoded_audio
-                    await pending.put(
-                        queue_item
-                    )
-                    queue_size = pending.qsize()
-                    break
-
-            if time.monotonic() > deadline:
-                return await handle_api_error(request, "job not initialized", 404)
-            await asyncio.sleep(0.05)
+        queue_size = await _enqueue_worker_items(
+            multi_job_id,
+            [
+                {
+                    "tensor": tensor,
+                    "worker_id": worker_id,
+                    "image_index": int(batch_idx),
+                    "is_last": is_last,
+                    "audio": decoded_audio,
+                }
+            ],
+        )
+        if queue_size is None:
+            return await handle_api_error(request, "job not initialized", 404)
 
         debug_log(
             f"job_complete received canonical envelope - job_id: {multi_job_id}, "

api/orchestration/__init__.py

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

api/orchestration/dispatch.py

@@ -1,33 +1,45 @@
 import asyncio
 import json
 import uuid
-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
-from ...utils.trace_logger import trace_debug, trace_info
-from ..schemas import coerce_positive_int
+try:
+    from ...utils.trace_logger import trace_debug, trace_info
+except ImportError:
+    def trace_debug(*_args, **_kwargs):
+        return None
+
+    def trace_info(*_args, **_kwargs):
+        return None
+
+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
 
 _least_busy_rr_index = 0
 
 
-async def worker_is_active(worker: dict[str, Any]) -> bool:
+async def worker_is_active(worker):
     """Ping worker's /prompt endpoint to confirm it's reachable."""
     url = build_worker_url(worker)
     return await probe_worker(url, timeout=3.0) is not None
 
 
-async def worker_ws_is_active(worker: dict[str, Any]) -> bool:
+async def worker_ws_is_active(worker):
     """Ping worker's websocket endpoint to confirm it's reachable."""
     session = await get_client_session()
     url = build_worker_url(worker, "/distributed/worker_ws")
-    headers = distributed_auth_headers(load_config())
     try:
-        ws = await session.ws_connect(url, heartbeat=20, timeout=3, headers=headers)
+        ws = await session.ws_connect(url, heartbeat=20, timeout=3)
         await ws.close()
         return True
     except asyncio.TimeoutError:
@@ -43,17 +55,16 @@ async def worker_ws_is_active(worker: dict[str, Any]) -> bool:
 
 async def _probe_worker_active(worker, use_websocket, semaphore):
     async with semaphore:
+        if worker.get("type") == "serverless":
+            from .serverless_runtime import get_manager
+            # Authentication/configuration failures must surface, not silently skip paid work.
+            status = await get_manager().health(worker)
+            return worker, status["ok"]
         is_active = await (worker_ws_is_active(worker) if use_websocket else worker_is_active(worker))
         return worker, is_active
 
 
-async def _dispatch_via_websocket(
-    worker_url,
-    payload,
-    client_id,
-    headers: dict[str, str],
-    timeout=60.0,
-):
+async def _dispatch_via_websocket(worker_url, payload, client_id, timeout=60.0):
     """Open a fresh worker websocket, dispatch one prompt, wait for ack, then close."""
     request_id = uuid.uuid4().hex
     ws_payload = {
@@ -67,12 +78,7 @@ async def _dispatch_via_websocket(
     ws_url = f"{ws_url}/distributed/worker_ws"
     session = await get_client_session()
 
-    async with session.ws_connect(
-        ws_url,
-        heartbeat=20,
-        timeout=timeout,
-        headers=headers,
-    ) as ws:
+    async with session.ws_connect(ws_url, heartbeat=20, timeout=timeout) as ws:
         await ws.send_json(ws_payload)
         async for msg in ws:
             if msg.type == aiohttp.WSMsgType.TEXT:
@@ -95,13 +101,13 @@ async def _dispatch_via_websocket(
 
 
 async def dispatch_worker_prompt(
-    worker: dict[str, Any],
-    prompt_obj: dict[str, Any],
-    workflow_meta: dict[str, Any] | None,
-    client_id: str | None = None,
-    use_websocket: bool = False,
-    trace_execution_id: str | None = None,
-) -> None:
+    worker,
+    prompt_obj,
+    workflow_meta,
+    client_id=None,
+    use_websocket=False,
+    trace_execution_id=None,
+):
     """Send the prepared prompt to a worker ComfyUI instance."""
     worker_url = build_worker_url(worker)
     url = build_worker_url(worker, "/prompt")
@@ -113,7 +119,6 @@ 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,
@@ -122,7 +127,6 @@ async def dispatch_worker_prompt(
                     "workflow": workflow_meta,
                 },
                 client_id,
-                ws_headers,
             )
             return
         except Exception as exc:
@@ -143,14 +147,14 @@ async def dispatch_worker_prompt(
 
 
 async def select_active_workers(
-    workers: list[dict[str, Any]],
-    use_websocket: bool,
-    delegate_master: bool,
-    trace_execution_id: str | None = None,
-    probe_concurrency: int = 8,
-) -> tuple[list[dict[str, Any]], bool]:
+    workers,
+    use_websocket,
+    delegate_master,
+    trace_execution_id=None,
+    probe_concurrency=8,
+):
     """Probe workers and return (active_workers, updated_delegate_master)."""
-    probe_limit = coerce_positive_int(probe_concurrency, 8)
+    probe_limit = parse_positive_int(probe_concurrency, 8)
     probe_semaphore = asyncio.Semaphore(probe_limit)
 
     if trace_execution_id and workers:
@@ -204,6 +208,10 @@ def _extract_queue_remaining(payload):
 
 async def _probe_worker_queue(worker, semaphore, probe_timeout):
     async with semaphore:
+        if worker.get("type") == "serverless":
+            from .serverless_runtime import get_manager
+            status = await get_manager().health(worker)
+            return {"worker": worker, "queue_remaining": status["queueRemaining"]}
         worker_url = build_worker_url(worker)
         payload = await probe_worker(worker_url, timeout=probe_timeout)
         if payload is None:
@@ -224,16 +232,16 @@ def _select_idle_round_robin(statuses):
 
 
 async def select_least_busy_worker(
-    workers: list[dict[str, Any]],
-    trace_execution_id: str | None = None,
-    probe_concurrency: int = 8,
-    probe_timeout: float = 3.0,
-) -> dict[str, Any] | None:
+    workers,
+    trace_execution_id=None,
+    probe_concurrency=8,
+    probe_timeout=3.0,
+):
     """Select one worker by queue depth, round-robin among idle workers."""
     if not workers:
         return None
 
-    probe_limit = coerce_positive_int(probe_concurrency, 8)
+    probe_limit = parse_positive_int(probe_concurrency, 8)
     probe_semaphore = asyncio.Semaphore(probe_limit)
     statuses = await asyncio.gather(
         *[
@@ -267,50 +275,3 @@ async def select_least_busy_worker(
             f"Least-busy worker selected: {worker.get('name')} ({worker.get('id')}), queue_remaining={queue_remaining}"
         )
     return worker
-
-
-async def rank_workers_by_load(
-    workers: list[dict[str, Any]],
-    trace_execution_id: str | None = None,
-    probe_concurrency: int = 8,
-    probe_timeout: float = 3.0,
-) -> list[dict[str, Any]]:
-    """Return all workers sorted by queue depth (ascending), preserving unreachable workers at the end."""
-    if not workers:
-        return []
-
-    probe_limit = coerce_positive_int(probe_concurrency, 8)
-    probe_semaphore = asyncio.Semaphore(probe_limit)
-    statuses = await asyncio.gather(
-        *[
-            _probe_worker_queue(worker, probe_semaphore, probe_timeout)
-            for worker in workers
-        ]
-    )
-
-    ranked_statuses = []
-    unreachable_workers = []
-    for worker, status in zip(workers, statuses):
-        if status is None:
-            unreachable_workers.append(worker)
-            continue
-        ranked_statuses.append(status)
-
-    ranked_statuses.sort(key=lambda status: status["queue_remaining"])
-    ranked_workers = [status["worker"] for status in ranked_statuses] + unreachable_workers
-
-    if trace_execution_id:
-        trace_debug(
-            trace_execution_id,
-            "Ranked workers by load: "
-            + ", ".join(
-                f"{status['worker'].get('id')}={status['queue_remaining']}"
-                for status in ranked_statuses
-            )
-            + (
-                f"; unreachable={','.join(worker.get('id', '') for worker in unreachable_workers)}"
-                if unreachable_workers
-                else ""
-            ),
-        )
-    return ranked_workers

api/orchestration/media_sync.py

@@ -3,12 +3,9 @@ import hashlib
 import mimetypes
 import os
 import re
-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
@@ -36,7 +33,7 @@ def _normalize_media_reference(value):
     return None
 
 
-def convert_paths_for_platform(obj: Any, target_separator: str) -> Any:
+def convert_paths_for_platform(obj, target_separator):
     """Recursively normalize likely file paths for the worker platform separator."""
     if target_separator not in ("/", "\\"):
         return obj
@@ -127,16 +124,12 @@ def _load_media_file_sync(filename):
     return file_bytes, file_hash, mime_type
 
 
-async def fetch_worker_path_separator(
-    worker: dict[str, Any],
-    trace_execution_id: str | None = None,
-) -> str | None:
+async def fetch_worker_path_separator(worker, trace_execution_id=None):
     """Best-effort fetch of a worker's path separator from /distributed/system_info."""
     url = build_worker_url(worker, "/distributed/system_info")
     session = await get_client_session()
-    headers = distributed_auth_headers(load_config())
     try:
-        async with session.get(url, headers=headers, timeout=aiohttp.ClientTimeout(total=5)) as resp:
+        async with session.get(url, timeout=aiohttp.ClientTimeout(total=5)) as resp:
             if resp.status != 200:
                 return None
             payload = await resp.json()
@@ -150,10 +143,33 @@ async def fetch_worker_path_separator(
         return None
 
 
+async def fetch_worker_has_node(worker, node_class, trace_execution_id=None):
+    """Best-effort check that a worker has ``node_class`` registered.
+
+    Uses ComfyUI's ``GET /object_info/{node_class}``, which returns ``{}`` for
+    unknown classes. Any failure counts as "not available" so callers fall
+    back to nodes every ComfyUI install has.
+    """
+    url = build_worker_url(worker, f"/object_info/{node_class}")
+    session = await get_client_session()
+    try:
+        async with session.get(url, timeout=aiohttp.ClientTimeout(total=5)) as resp:
+            if resp.status != 200:
+                return False
+            payload = await resp.json()
+            return isinstance(payload, dict) and node_class in payload
+    except Exception as exc:
+        message = f"Failed to probe node '{node_class}' on worker {worker.get('id')}: {exc}"
+        if trace_execution_id:
+            trace_debug(trace_execution_id, message)
+        else:
+            debug_log(f"[Distributed] {message}")
+        return False
+
+
 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")
@@ -161,7 +177,6 @@ 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:
@@ -202,11 +217,7 @@ async def _upload_media_to_worker(worker, filename, file_bytes, file_hash, mime_
     return True, worker_path
 
 
-async def sync_worker_media(
-    worker: dict[str, Any],
-    prompt_obj: dict[str, Any],
-    trace_execution_id: str | None = None,
-) -> None:
+async def sync_worker_media(worker, prompt_obj, trace_execution_id=None):
     """Sync referenced media files from master to a remote worker before dispatch."""
     media_refs = _find_media_references(prompt_obj)
     if not media_refs:

api/orchestration/prompt_transform.py

@@ -1,6 +1,5 @@
 import json
 from collections import deque
-from typing import Any
 
 from ...utils.logging import debug_log
 
@@ -28,7 +27,7 @@ class PromptIndex:
     def nodes_for_class(self, class_name):
         return self.nodes_by_class.get(class_name, [])
 
-    def has_upstream(self, start_node_id: str, target_class: str) -> bool:
+    def has_upstream(self, start_node_id, target_class):
         cache_key = (str(start_node_id), target_class)
         if cache_key in self._upstream_cache:
             return self._upstream_cache[cache_key]
@@ -89,38 +88,6 @@ def _find_downstream_nodes(prompt_obj, start_ids):
     return connected
 
 
-def _find_downstream_of_output_slot(prompt_obj, node_id, slot_index):
-    """Return all nodes reachable from a specific output slot on one node."""
-    source_id = str(node_id)
-    try:
-        slot = int(slot_index)
-    except (TypeError, ValueError):
-        return set()
-
-    direct_consumers = set()
-    for candidate_id, candidate_node in _iter_prompt_nodes(prompt_obj):
-        inputs = candidate_node.get("inputs", {})
-        for value in inputs.values():
-            if not (isinstance(value, list) and len(value) == 2):
-                continue
-            if str(value[0]) != source_id:
-                continue
-            try:
-                input_slot = int(value[1])
-            except (TypeError, ValueError):
-                debug_log(
-                    f"Prompt transform: skipping malformed slot reference for node {candidate_id}: {value}"
-                )
-                continue
-            if input_slot == slot:
-                direct_consumers.add(str(candidate_id))
-                break
-
-    if not direct_consumers:
-        return set()
-    return _find_downstream_nodes(prompt_obj, list(direct_consumers))
-
-
 def _create_numeric_id_generator(prompt_obj):
     """Return a closure that yields new numeric string IDs."""
     max_id = 0
@@ -128,7 +95,6 @@ def _create_numeric_id_generator(prompt_obj):
         try:
             numeric = int(node_id)
         except (TypeError, ValueError):
-            debug_log(f"Prompt transform: ignoring non-numeric node id while generating ids: {node_id}")
             continue
         max_id = max(max_id, numeric)
 
@@ -159,193 +125,229 @@ def _find_upstream_nodes(prompt_obj, start_ids):
     return connected
 
 
-def _resolve_participants(enabled_worker_ids, delegate_master):
-    worker_ids = [str(worker_id) for worker_id in (enabled_worker_ids or [])]
-    if delegate_master:
-        return worker_ids
-    return ["master"] + worker_ids
+_DELEGATE_MASTER_RETAINED_UPSTREAM_CLASSES = {
+    "PrimitiveBoolean",
+    "PrimitiveFloat",
+    "PrimitiveInt",
+    "PrimitiveNode",
+    "PrimitiveString",
+}
+
+_DELEGATE_MASTER_ALWAYS_RETAINED_UPSTREAM_CLASSES = {
+    "LoadImage",
+}
+
+_DELEGATE_MASTER_SAFE_SCALAR_TYPES = {"BOOLEAN", "FLOAT", "INT", "STRING"}
+_DELEGATE_MASTER_SAFE_LIST_TYPES = {"LIST"}
+
+# ComfyUI 0.23 exposes CreateList via the newer schema API rather than the
+# legacy RETURN_TYPES/INPUT_TYPES attributes. Treat it as a safe config utility
+# only after its connected inputs recursively prove safe.
+_DELEGATE_MASTER_SAFE_DYNAMIC_CONFIG_OUTPUT_CLASSES = {"CreateList"}
+_DELEGATE_MASTER_SAFE_DYNAMIC_CONFIG_INPUT_PREFIXES = {
+    "CreateList": ("inputs.",),
+}
+
+# Test hook. At runtime this stays None and the ComfyUI node registry is loaded lazily.
+_DELEGATE_MASTER_NODE_CLASS_MAPPINGS = None
 
 
-def _remove_dangling_input_refs(prompt_obj):
-    """Drop input links that point to nodes no longer present in prompt_obj."""
-    existing_ids = set(prompt_obj.keys())
-    for _node_id, node in _iter_prompt_nodes(prompt_obj):
-        inputs = node.get("inputs", {})
-        for input_name, input_value in list(inputs.items()):
-            if isinstance(input_value, list) and len(input_value) == 2:
-                source_id = str(input_value[0])
-                if source_id not in existing_ids:
-                    inputs.pop(input_name, None)
+def _get_delegate_master_node_class_mappings():
+    """Return ComfyUI node-class mappings when available."""
+    if _DELEGATE_MASTER_NODE_CLASS_MAPPINGS is not None:
+        return _DELEGATE_MASTER_NODE_CLASS_MAPPINGS
+    try:
+        import nodes as comfy_nodes  # type: ignore
+    except Exception:  # pragma: no cover - depends on ComfyUI runtime imports
+        return {}
+    return getattr(comfy_nodes, "NODE_CLASS_MAPPINGS", {}) or {}
 
 
-def _has_terminal_output_nodes(prompt_obj):
-    """Return True when the prompt already has at least one terminal output node."""
-    for _node_id, node in _iter_prompt_nodes(prompt_obj):
-        if node.get("class_type") in {"PreviewImage", "SaveImage"}:
+def _get_delegate_master_node_class(class_type):
+    mappings = _get_delegate_master_node_class_mappings()
+    return mappings.get(class_type) if isinstance(mappings, dict) else None
+
+
+def _normalize_delegate_master_return_type(return_type):
+    if return_type is None:
+        return ""
+    return str(return_type).strip().upper()
+
+
+def _delegate_master_type_is_safe_scalar(type_name):
+    return type_name in _DELEGATE_MASTER_SAFE_SCALAR_TYPES
+
+
+def _delegate_master_type_is_safe_config(type_name):
+    return _delegate_master_type_is_safe_scalar(type_name) or type_name in _DELEGATE_MASTER_SAFE_LIST_TYPES
+
+
+def _delegate_master_output_is_safe_scalar(class_type, output_index):
+    """Return True when a registered node output is lightweight config data."""
+    if class_type in _DELEGATE_MASTER_SAFE_DYNAMIC_CONFIG_OUTPUT_CLASSES:
+        return True
+    node_class = _get_delegate_master_node_class(class_type)
+    return_types = getattr(node_class, "RETURN_TYPES", ()) if node_class is not None else ()
+    try:
+        output_type = return_types[int(output_index)]
+    except (IndexError, TypeError, ValueError):
+        return False
+    return _delegate_master_type_is_safe_config(_normalize_delegate_master_return_type(output_type))
+
+
+def _get_delegate_master_input_types(class_type):
+    node_class = _get_delegate_master_node_class(class_type)
+    input_types = getattr(node_class, "INPUT_TYPES", None) if node_class is not None else None
+    if callable(input_types):
+        try:
+            input_types = input_types()
+        except TypeError:
+            return {}
+    return input_types if isinstance(input_types, dict) else {}
+
+
+def _normalize_delegate_master_input_type(input_spec):
+    if isinstance(input_spec, (list, tuple)) and input_spec:
+        return _normalize_delegate_master_return_type(input_spec[0])
+    return _normalize_delegate_master_return_type(input_spec)
+
+
+def _delegate_master_input_is_safe_scalar(class_type, input_name):
+    """Return True when a registered downstream input expects config data."""
+    for prefix in _DELEGATE_MASTER_SAFE_DYNAMIC_CONFIG_INPUT_PREFIXES.get(class_type, ()):
+        if input_name.startswith(prefix):
             return True
+    input_types = _get_delegate_master_input_types(class_type)
+    for section_name in ("required", "optional"):
+        section = input_types.get(section_name, {})
+        if isinstance(section, dict) and input_name in section:
+            input_type = _normalize_delegate_master_input_type(section[input_name])
+            return _delegate_master_type_is_safe_config(input_type)
     return False
 
 
-def _add_preview_node(prompt_obj, next_id_fn, source_node_id, slot_index, title_suffix):
-    """Attach a PreviewImage output node to the given source output slot."""
-    preview_id = next_id_fn()
-    prompt_obj[preview_id] = {
-        "inputs": {
-            "images": [str(source_node_id), int(slot_index)],
-        },
-        "class_type": "PreviewImage",
-        "_meta": {
-            "title": f"Preview Image ({title_suffix})",
-        },
-    }
-    return preview_id
+def _is_delegate_master_always_retained_upstream_node(node):
+    if not isinstance(node, dict):
+        return False
+    class_type = node.get("class_type")
+    return isinstance(class_type, str) and class_type in _DELEGATE_MASTER_ALWAYS_RETAINED_UPSTREAM_CLASSES
 
 
-def _ensure_worker_output_node(prompt_obj):
-    """Guarantee worker prompts retain at least one output node after pruning.
+def _is_delegate_master_retained_upstream_node(node, output_index=0):
+    """Return True for lightweight upstream nodes safe to keep on the master."""
+    if not isinstance(node, dict):
+        return False
+    class_type = node.get("class_type")
+    if not isinstance(class_type, str):
+        return False
+    return (
+        class_type in _DELEGATE_MASTER_RETAINED_UPSTREAM_CLASSES
+        or class_type.startswith("Primitive")
+        or _delegate_master_output_is_safe_scalar(class_type, output_index)
+    )
 
-    Branch pruning can remove all explicit output nodes for worker prompts, which
-    causes ComfyUI validation to reject the prompt (`prompt_no_outputs`).
-    """
-    if _has_terminal_output_nodes(prompt_obj):
-        return prompt_obj
 
-    next_id = _create_numeric_id_generator(prompt_obj)
+def _collect_delegate_master_retained_upstream_branch(
+    prompt_obj,
+    node_id,
+    output_index,
+    memo,
+    visiting,
+):
+    """Return safe retained branch nodes, or None when the branch is not safe."""
+    node_id = str(node_id)
+    cache_key = (node_id, output_index)
+    if cache_key in memo:
+        cached = memo[cache_key]
+        return None if cached is None else set(cached)
+    if cache_key in visiting:
+        memo[cache_key] = None
+        return None
 
-    # Prefer a DistributedBranchCollector output tied to the assigned branch for this worker.
-    for node_id, node in _iter_prompt_nodes(prompt_obj):
-        if node.get("class_type") != "DistributedBranchCollector":
+    node = prompt_obj.get(node_id)
+    if not _is_delegate_master_retained_upstream_node(node, output_index):
+        memo[cache_key] = None
+        return None
+
+    visiting.add(cache_key)
+    retained = {node_id}
+    inputs = node.get("inputs", {}) if isinstance(node, dict) else {}
+    class_type = node.get("class_type") if isinstance(node, dict) else None
+    for input_name, value in inputs.items():
+        if not (isinstance(value, list) and len(value) == 2):
             continue
+        if not _delegate_master_input_is_safe_scalar(class_type, input_name):
+            visiting.remove(cache_key)
+            memo[cache_key] = None
+            return None
+        source_id = str(value[0])
+        branch = _collect_delegate_master_retained_upstream_branch(
+            prompt_obj,
+            source_id,
+            value[1],
+            memo,
+            visiting,
+        )
+        if branch is None:
+            visiting.remove(cache_key)
+            memo[cache_key] = None
+            return None
+        retained.update(branch)
+
+    visiting.remove(cache_key)
+    memo[cache_key] = frozenset(retained)
+    return retained
+
+
+def _find_delegate_master_retained_upstream_nodes(prompt_obj, start_ids):
+    """Return lightweight upstream nodes needed by kept delegate-master nodes."""
+    connected = set()
+    memo = {}
+    for node_id in start_ids:
+        node = prompt_obj.get(str(node_id)) or {}
         inputs = node.get("inputs", {})
-        try:
-            assigned_branch = int(inputs.get("assigned_branch", -1))
-        except (TypeError, ValueError):
-            assigned_branch = -1
-        if assigned_branch >= 0:
-            _add_preview_node(prompt_obj, next_id, node_id, assigned_branch, "auto-added worker output")
-            return prompt_obj
-
-    # Otherwise, anchor to the assigned DistributedBranch slot if available.
-    for node_id, node in _iter_prompt_nodes(prompt_obj):
-        if node.get("class_type") != "DistributedBranch":
-            continue
-        inputs = node.get("inputs", {})
-        try:
-            assigned_branch = int(inputs.get("assigned_branch", -1))
-        except (TypeError, ValueError):
-            assigned_branch = -1
-        if assigned_branch >= 0:
-            _add_preview_node(prompt_obj, next_id, node_id, assigned_branch, "auto-added worker output")
-            return prompt_obj
-
-    # Idle participants (assigned_branch=-1) still need a terminal node to pass
-    # validation; keep this cheap by emitting a tiny synthetic image preview.
-    empty_id = next_id()
-    prompt_obj[empty_id] = {
-        "class_type": "DistributedEmptyImage",
-        "inputs": {
-            "height": 64,
-            "width": 64,
-            "channels": 3,
-        },
-        "_meta": {
-            "title": "Distributed Empty Image (auto-added worker output)",
-        },
-    }
-    _add_preview_node(prompt_obj, next_id, empty_id, 0, "auto-added worker output")
-    return prompt_obj
-
-
-def _prune_worker_downstream_of_branch_collectors(prompt_obj):
-    """Drop worker-side nodes downstream of branch collectors."""
-    collector_ids = find_nodes_by_class(prompt_obj, "DistributedBranchCollector")
-    if not collector_ids:
-        return prompt_obj
-
-    downstream = _find_downstream_nodes(prompt_obj, collector_ids)
-    for node_id in downstream:
-        if node_id in collector_ids:
-            continue
-        prompt_obj.pop(node_id, None)
-
-    _remove_dangling_input_refs(prompt_obj)
-    return prompt_obj
-
-
-def prune_prompt_for_branch_worker(
-    prompt_obj: dict[str, Any],
-    branch_node_id: str,
-    assigned_branch: int | list[int] | tuple[int, ...] | set[int],
-    num_branches: int,
-) -> dict[str, Any]:
-    """Prune non-assigned branch paths while keeping shared downstream nodes."""
-    branch_id = str(branch_node_id)
-
-    assigned_slots = set()
-    if isinstance(assigned_branch, (list, tuple, set)):
-        for value in assigned_branch:
-            try:
-                idx = int(value)
-            except (TypeError, ValueError):
-                debug_log(f"Prompt transform: invalid assigned branch slot ignored: {value}")
+        class_type = node.get("class_type") if isinstance(node, dict) else None
+        for input_name, value in inputs.items():
+            if not (isinstance(value, list) and len(value) == 2):
                 continue
-            if idx >= 0:
-                assigned_slots.add(idx)
-    else:
-        try:
-            idx = int(assigned_branch)
-        except (TypeError, ValueError):
-            idx = -1
-        if idx >= 0:
-            assigned_slots.add(idx)
-
-    try:
-        total_slots = int(num_branches)
-    except (TypeError, ValueError):
-        total_slots = 2
-    total_slots = max(2, min(total_slots, 10))
-
-    downstream_by_slot = {}
-    for slot_idx in range(total_slots):
-        downstream = _find_downstream_of_output_slot(prompt_obj, branch_id, slot_idx)
-        downstream.discard(branch_id)
-        downstream_by_slot[slot_idx] = downstream
-
-    keep = {branch_id}
-    for slot_idx in assigned_slots:
-        keep.update(downstream_by_slot.get(slot_idx, set()))
-
-    remove = set()
-    for slot_idx in range(total_slots):
-        if slot_idx in assigned_slots:
-            continue
-        remove.update(downstream_by_slot.get(slot_idx, set()))
-
-    remove -= keep
-    for node_id in remove:
-        prompt_obj.pop(node_id, None)
-
-    _remove_dangling_input_refs(prompt_obj)
-    return prompt_obj
+            source_node = prompt_obj.get(str(value[0]))
+            if _is_delegate_master_always_retained_upstream_node(source_node):
+                connected.add(str(value[0]))
+                continue
+            if not _delegate_master_input_is_safe_scalar(class_type, input_name):
+                continue
+            branch = _collect_delegate_master_retained_upstream_branch(
+                prompt_obj,
+                value[0],
+                value[1],
+                memo,
+                set(),
+            )
+            if branch is not None:
+                connected.update(branch)
+    return connected
 
 
-def prune_prompt_for_worker(prompt_obj: dict[str, Any]) -> dict[str, Any]:
-    """Prune worker prompt to distributed nodes and their upstream dependencies."""
+WORKER_SINK_CLASS = "DistributedWorkerSink"
+
+
+def prune_prompt_for_worker(prompt_obj, use_worker_sink=False):
+    """Prune worker prompt to distributed nodes and their upstream dependencies.
+
+    Removing downstream nodes can leave the worker prompt without an output
+    node. One is re-added after each such distributed node: the hidden
+    ``DistributedWorkerSink`` when ``use_worker_sink`` is set (the worker must
+    have it registered), otherwise PreviewImage/PreviewAudio for workers
+    running older versions. The preview nodes write every frame to the
+    worker's temp directory, which the sink avoids.
+    """
     collector_ids = find_nodes_by_class(prompt_obj, "DistributedCollector")
-    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 + branch_collector_ids + upscale_ids + branch_ids
+    distributed_ids = collector_ids + upscale_ids
     if not distributed_ids:
         return prompt_obj
 
     connected = _find_upstream_nodes(prompt_obj, distributed_ids)
-    if branch_ids:
-        connected.update(_find_downstream_nodes(prompt_obj, branch_ids))
-    if branch_collector_ids:
-        downstream_of_collectors = _find_downstream_nodes(prompt_obj, branch_collector_ids)
-        connected -= (downstream_of_collectors - set(branch_collector_ids))
-
     pruned_prompt = {}
     for node_id in connected:
         node = prompt_obj.get(node_id)
@@ -358,30 +360,53 @@ def prune_prompt_for_worker(prompt_obj: dict[str, Any]) -> dict[str, Any]:
         if dist_id not in pruned_prompt:
             continue
         downstream = _find_downstream_nodes(prompt_obj, [dist_id])
-        has_removed_downstream = any(node_id != dist_id and node_id not in connected for node_id in downstream)
+        has_removed_downstream = any(node_id != dist_id for node_id in downstream)
         if has_removed_downstream:
-            preview_id = next_id()
-            pruned_prompt[preview_id] = {
-                "inputs": {
-                    "images": [dist_id, 0],
-                },
-                "class_type": "PreviewImage",
-                "_meta": {
-                    "title": "Preview Image (auto-added)",
-                },
-            }
+            original_node = prompt_obj.get(str(dist_id), {})
+            class_type = original_node.get("class_type")
+            inputs = original_node.get("inputs", {})
+            image_connected = class_type != "DistributedCollector" or (
+                isinstance(inputs.get("images"), list)
+                and len(inputs["images"]) == 2
+            )
+            audio_connected = (
+                class_type == "DistributedCollector"
+                and isinstance(inputs.get("audio"), list)
+                and len(inputs["audio"]) == 2
+            )
+
+            if use_worker_sink and (image_connected or audio_connected):
+                pruned_prompt[next_id()] = {
+                    "inputs": {"value": [dist_id, 0 if image_connected else 1]},
+                    "class_type": WORKER_SINK_CLASS,
+                    "_meta": {"title": "Distributed Worker Sink (auto-added)"},
+                }
+            elif image_connected:
+                preview_id = next_id()
+                pruned_prompt[preview_id] = {
+                    "inputs": {"images": [dist_id, 0]},
+                    "class_type": "PreviewImage",
+                    "_meta": {"title": "Preview Image (auto-added)"},
+                }
+            elif audio_connected:
+                preview_id = next_id()
+                pruned_prompt[preview_id] = {
+                    "inputs": {"audio": [dist_id, 1]},
+                    "class_type": "PreviewAudio",
+                    "_meta": {"title": "Preview Audio (auto-added)"},
+                }
 
     return pruned_prompt
 
 
-def prepare_delegate_master_prompt(
-    prompt_obj: dict[str, Any],
-    collector_ids: list[str],
-) -> dict[str, Any]:
+def prepare_delegate_master_prompt(prompt_obj, collector_ids):
     """Prune master prompt so it only executes post-collector nodes in delegate mode."""
     downstream = _find_downstream_nodes(prompt_obj, collector_ids)
     nodes_to_keep = set(collector_ids)
     nodes_to_keep.update(downstream)
+    nodes_to_keep.update(
+        _find_delegate_master_retained_upstream_nodes(prompt_obj, nodes_to_keep)
+    )
 
     pruned_prompt = {}
     for node_id in nodes_to_keep:
@@ -409,16 +434,9 @@ def prepare_delegate_master_prompt(
         collector_entry = pruned_prompt.get(collector_id)
         if not collector_entry:
             continue
-        collector_class = collector_entry.get("class_type")
-        if collector_class == "DistributedBranchCollector":
-            # Branch collector does not require an image input in delegate-only mode.
-            continue
-        if collector_class == "UltimateSDUpscaleDistributed":
-            # Swap to lightweight delegate collector — no model inputs needed.
-            collector_entry["class_type"] = "USDUDelegateCollector"
-            debug_log(
-                f"Swapped USDU node {collector_id} to USDUDelegateCollector for delegate-only master prompt."
-            )
+        original_inputs = (prompt_obj.get(collector_id) or {}).get("inputs", {})
+        original_images = original_inputs.get("images")
+        if not (isinstance(original_images, list) and len(original_images) == 2):
             continue
         placeholder_id = next_id()
         pruned_prompt[placeholder_id] = {
@@ -440,145 +458,32 @@ def prepare_delegate_master_prompt(
     return pruned_prompt
 
 
-def generate_job_id_map(prompt_index: PromptIndex, prefix: str) -> dict[str, str]:
+def generate_job_id_map(prompt_index, prefix):
     """Create stable per-node job IDs for distributed nodes."""
     job_map = {}
-    distributed_nodes = (
-        prompt_index.nodes_for_class("DistributedCollector")
-        + prompt_index.nodes_for_class("DistributedBranchCollector")
-        + prompt_index.nodes_for_class("UltimateSDUpscaleDistributed")
-        + prompt_index.nodes_for_class("DistributedBranch")
+    distributed_nodes = prompt_index.nodes_for_class("DistributedCollector") + prompt_index.nodes_for_class(
+        "UltimateSDUpscaleDistributed"
     )
     for node_id in distributed_nodes:
         job_map[node_id] = f"{prefix}_{node_id}"
     return job_map
 
 
-def _override_simple_nodes(prompt_copy, prompt_index, is_master, participant_id, enabled_json, class_names):
-    """Configure simple distributed nodes (DistributedSeed, DistributedValue)."""
-    for class_name in class_names:
-        for node_id in prompt_index.nodes_for_class(class_name):
-            node = prompt_copy.get(node_id)
-            if not isinstance(node, dict):
-                continue
-            inputs = node.setdefault("inputs", {})
-            inputs["is_worker"] = not is_master
-            inputs["enabled_worker_ids"] = enabled_json
-            inputs["worker_id"] = "" if is_master else str(participant_id)
-
-
-def _override_job_nodes(
-    prompt_copy, prompt_index, is_master, participant_id,
-    job_id_map, master_url, enabled_json, delegate_master,
-    class_names, skip_if_upstream=None,
-):
-    """Configure job-aware distributed nodes (collector, upscale)."""
-    for class_name in class_names:
-        for node_id in prompt_index.nodes_for_class(class_name):
-            node = prompt_copy.get(node_id)
-            if not isinstance(node, dict):
-                continue
-
-            if skip_if_upstream and prompt_index.has_upstream(node_id, skip_if_upstream):
-                node.setdefault("inputs", {})["pass_through"] = True
-                continue
-
-            inputs = node.setdefault("inputs", {})
-            inputs["multi_job_id"] = job_id_map.get(node_id, node_id)
-            inputs["is_worker"] = not is_master
-            inputs["enabled_worker_ids"] = enabled_json
-            if is_master:
-                inputs["delegate_only"] = bool(delegate_master)
-                inputs.pop("master_url", None)
-                inputs.pop("worker_id", None)
-            else:
-                inputs["master_url"] = master_url
-                inputs["worker_id"] = participant_id
-                inputs["delegate_only"] = False
-
-
-def _override_branch_nodes(
-    prompt_copy,
-    prompt_index,
-    is_master,
-    participant_id,
-    enabled_worker_ids,
-    job_id_map,
-    delegate_master,
-):
-    """Assign branch slots per participant and prune non-assigned branch paths."""
-    participants = _resolve_participants(enabled_worker_ids, delegate_master)
-    participant_id = str(participant_id)
-    participant_count = len(participants)
-    participant_pos = participants.index(participant_id) if participant_id in participants else None
-
-    for node_id in prompt_index.nodes_for_class("DistributedBranch"):
+def _override_seed_nodes(prompt_copy, prompt_index, is_master, participant_id, worker_index_map):
+    """Configure DistributedSeed nodes for master or worker role."""
+    for node_id in prompt_index.nodes_for_class("DistributedSeed"):
         node = prompt_copy.get(node_id)
         if not isinstance(node, dict):
             continue
-
         inputs = node.setdefault("inputs", {})
-        try:
-            num_branches = int(inputs.get("num_branches", 2))
-        except (TypeError, ValueError):
-            num_branches = 2
-        num_branches = max(2, min(num_branches, 10))
-
-        assigned_slots = []
-        if participant_pos is not None and participant_count > 0:
-            assigned_slots = [slot for slot in range(num_branches) if slot % participant_count == participant_pos]
-
-        if len(assigned_slots) == 1:
-            assigned_branch = assigned_slots[0]
-        else:
-            assigned_branch = -1
-
         inputs["is_worker"] = not is_master
-        inputs["worker_id"] = "" if is_master else participant_id
-        inputs["assigned_branch"] = assigned_branch
-        inputs["multi_job_id"] = job_id_map.get(node_id, node_id)
-
-        if is_master and delegate_master:
-            continue
-
-        prompt_copy = prune_prompt_for_branch_worker(
-            prompt_copy,
-            node_id,
-            assigned_slots,
-            num_branches,
-        )
-
-    return prompt_copy
+        if is_master:
+            inputs["worker_id"] = ""
+        else:
+            inputs["worker_id"] = f"worker_{worker_index_map.get(participant_id, 0)}"
 
 
-def _find_upstream_branch_node(prompt_copy, start_node_id):
-    """Locate the nearest upstream DistributedBranch node for the provided node."""
-    visited = set()
-    stack = [str(start_node_id)]
-    while stack:
-        node_id = stack.pop()
-        if node_id in visited:
-            continue
-        visited.add(node_id)
-        node = prompt_copy.get(node_id)
-        if not isinstance(node, dict):
-            continue
-
-        inputs = node.get("inputs", {})
-        for input_value in inputs.values():
-            if not (isinstance(input_value, list) and len(input_value) == 2):
-                continue
-            source_id = str(input_value[0])
-            source_node = prompt_copy.get(source_id)
-            if not isinstance(source_node, dict):
-                continue
-            if source_node.get("class_type") == "DistributedBranch":
-                return source_id
-            stack.append(source_id)
-    return None
-
-
-def _override_branch_collector_nodes(
+def _override_collector_nodes(
     prompt_copy,
     prompt_index,
     is_master,
@@ -588,77 +493,87 @@ def _override_branch_collector_nodes(
     enabled_json,
     delegate_master,
 ):
-    """Configure DistributedBranchCollector nodes for branch convergence."""
-    for node_id in prompt_index.nodes_for_class("DistributedBranchCollector"):
+    """Configure DistributedCollector nodes for master or worker role."""
+    for node_id in prompt_index.nodes_for_class("DistributedCollector"):
         node = prompt_copy.get(node_id)
         if not isinstance(node, dict):
             continue
 
-        upstream_branch_id = _find_upstream_branch_node(prompt_copy, node_id)
-        assigned_branch = -1
-        if upstream_branch_id is not None:
-            branch_node = prompt_copy.get(upstream_branch_id, {})
-            branch_inputs = branch_node.get("inputs", {}) if isinstance(branch_node, dict) else {}
-            try:
-                assigned_branch = int(branch_inputs.get("assigned_branch", -1))
-            except (TypeError, ValueError):
-                assigned_branch = -1
+        if prompt_index.has_upstream(node_id, "UltimateSDUpscaleDistributed"):
+            node.setdefault("inputs", {})["pass_through"] = True
+            continue
 
         inputs = node.setdefault("inputs", {})
-        # Group all branch collectors under the same DistributedBranch into one queue.
-        job_source_id = upstream_branch_id or node_id
-        inputs["multi_job_id"] = job_id_map.get(job_source_id, job_source_id)
+        inputs["multi_job_id"] = job_id_map.get(node_id, node_id)
         inputs["is_worker"] = not is_master
         inputs["enabled_worker_ids"] = enabled_json
-        inputs["assigned_branch"] = assigned_branch
         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["worker_id"] = participant_id
             inputs["delegate_only"] = False
 
 
+def _override_upscale_nodes(
+    prompt_copy,
+    prompt_index,
+    is_master,
+    participant_id,
+    job_id_map,
+    master_url,
+    enabled_json,
+):
+    """Configure UltimateSDUpscaleDistributed nodes for master or worker role."""
+    for node_id in prompt_index.nodes_for_class("UltimateSDUpscaleDistributed"):
+        node = prompt_copy.get(node_id)
+        if not isinstance(node, dict):
+            continue
+        inputs = node.setdefault("inputs", {})
+        inputs["multi_job_id"] = job_id_map.get(node_id, node_id)
+        inputs["is_worker"] = not is_master
+        inputs["enabled_worker_ids"] = enabled_json
+        if is_master:
+            inputs.pop("master_url", None)
+            inputs.pop("worker_id", None)
+        else:
+            inputs["master_url"] = master_url
+            inputs["worker_id"] = participant_id
+
+
+def _override_value_nodes(prompt_copy, prompt_index, is_master, participant_id, worker_index_map):
+    """Configure DistributedValue nodes for master or worker role."""
+    for node_id in prompt_index.nodes_for_class("DistributedValue"):
+        node = prompt_copy.get(node_id)
+        if not isinstance(node, dict):
+            continue
+        inputs = node.setdefault("inputs", {})
+        inputs["is_worker"] = not is_master
+        if is_master:
+            inputs["worker_id"] = ""
+        else:
+            inputs["worker_id"] = f"worker_{worker_index_map.get(participant_id, 0)}"
+
+
 def apply_participant_overrides(
-    prompt_copy: dict[str, Any],
-    participant_id: str,
-    enabled_worker_ids: list[str],
-    job_id_map: dict[str, str],
-    master_url: str,
-    delegate_master: bool,
-    prompt_index: PromptIndex,
-) -> dict[str, Any]:
+    prompt_copy,
+    participant_id,
+    enabled_worker_ids,
+    job_id_map,
+    master_url,
+    delegate_master,
+    prompt_index,
+):
     """Return a prompt copy with hidden inputs configured for master/worker."""
     is_master = participant_id == "master"
+    worker_index_map = {wid: idx for idx, wid in enumerate(enabled_worker_ids)}
     enabled_json = json.dumps(enabled_worker_ids)
 
-    _override_simple_nodes(
-        prompt_copy, prompt_index, is_master, participant_id, enabled_json,
-        ["DistributedSeed", "DistributedValue"],
-    )
-    _override_job_nodes(
-        prompt_copy, prompt_index, is_master, participant_id,
-        job_id_map, master_url, enabled_json, delegate_master,
-        ["DistributedCollector"],
-        skip_if_upstream="UltimateSDUpscaleDistributed",
-    )
-    _override_job_nodes(
-        prompt_copy, prompt_index, is_master, participant_id,
-        job_id_map, master_url, enabled_json, delegate_master,
-        ["UltimateSDUpscaleDistributed"],
-    )
-    prompt_copy = _override_branch_nodes(
-        prompt_copy,
-        prompt_index,
-        is_master,
-        participant_id,
-        enabled_worker_ids,
-        job_id_map,
-        delegate_master,
-    )
-    _override_branch_collector_nodes(
+    _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_collector_nodes(
         prompt_copy,
         prompt_index,
         is_master,
@@ -668,9 +583,14 @@ def apply_participant_overrides(
         enabled_json,
         delegate_master,
     )
-
-    if not is_master:
-        prompt_copy = _prune_worker_downstream_of_branch_collectors(prompt_copy)
-        prompt_copy = _ensure_worker_output_node(prompt_copy)
+    _override_upscale_nodes(
+        prompt_copy,
+        prompt_index,
+        is_master,
+        participant_id,
+        job_id_map,
+        master_url,
+        enabled_json,
+    )
 
     return prompt_copy

api/orchestration/runpod_client.py

@@ -0,0 +1,123 @@
+"""Small Runpod queue API client. Never retries a potentially billable POST."""
+import asyncio
+import json
+import re
+
+import aiohttp
+
+API_BASE = "https://api.runpod.ai/v2"
+MAX_REQUEST_BYTES = 10_000_000
+_ID = re.compile(r"[A-Za-z0-9_-]{1,200}\Z")
+
+
+class RunpodError(RuntimeError):
+    """Sanitized provider failure safe to display in the panel."""
+
+
+class AmbiguousSubmission(RunpodError):
+    """Submission may have been accepted: never retry it automatically."""
+
+
+def serialize_request(workflow, images, timeout_seconds):
+    """Serialize and size-check the exact JSON body used by Runpod ``/run``."""
+    payload = {
+        "input": {"workflow": workflow, "images": images},
+        "policy": {
+            "executionTimeout": timeout_seconds * 1000,
+            "ttl": timeout_seconds * 1000,
+        },
+    }
+    body = json.dumps(
+        payload,
+        allow_nan=False,
+        ensure_ascii=False,
+        separators=(",", ":"),
+    ).encode("utf-8")
+    size = len(body)
+    if size > MAX_REQUEST_BYTES:
+        raise ValueError(
+            f"Serialized serverless request body is {size:,} bytes; Runpod allows at most "
+            f"{MAX_REQUEST_BYTES:,} bytes. Base64-encoded images add about 33% over raw "
+            "image bytes; reduce image or workflow size."
+        )
+    return body
+
+
+def checked_id(value):
+    if not isinstance(value, str) or not _ID.fullmatch(value):
+        raise ValueError("Invalid Runpod endpoint or job ID")
+    return value
+
+
+class RunpodClient:
+    def __init__(self, session, endpoint_id, api_key):
+        self.session = session
+        self.endpoint_id = checked_id(endpoint_id)
+        if not isinstance(api_key, str) or not api_key or any(c in api_key for c in "\r\n"):
+            raise ValueError("A valid Runpod API key is required")
+        self._api_key = api_key
+
+    async def _request(self, method, path, payload=None, request_body=None):
+        # Base origin is code-owned, never a user-supplied URL. Do not follow redirects
+        # with credentials, or include response bodies/transport errors in exceptions.
+        url = f"{API_BASE}/{self.endpoint_id}/{path}"
+        attempts = 3 if method == "GET" else 1
+        for attempt in range(attempts):
+            try:
+                request_kwargs = {
+                    "headers": {"Authorization": f"Bearer {self._api_key}", "Content-Type": "application/json"},
+                    "timeout": aiohttp.ClientTimeout(total=30),
+                    "allow_redirects": False,
+                }
+                if request_body is None:
+                    request_kwargs["json"] = payload
+                else:
+                    request_kwargs["data"] = request_body
+                async with self.session.request(
+                    method, url,
+                    **request_kwargs,
+                ) as response:
+                    status = response.status
+                    if status != 200:
+                        if method == "GET" and (status == 429 or status >= 500) and attempt + 1 < attempts:
+                            await asyncio.sleep(0.5 * (2 ** attempt))
+                            continue
+                        if method == "POST" and path == "run" and status >= 500:
+                            raise AmbiguousSubmission("Runpod submission returned a server error; acceptance is unknown. Check the provider before retrying.")
+                        raise RunpodError(f"Runpod {path.split('/')[0]} request failed (HTTP {status})")
+                    try:
+                        result = await response.json()
+                    except (ValueError, aiohttp.ContentTypeError):
+                        if method == "POST" and path == "run":
+                            raise AmbiguousSubmission("Runpod submission returned an unreadable response; check the provider before retrying.") from None
+                        raise RunpodError("Runpod returned an invalid JSON response") from None
+                    if not isinstance(result, dict):
+                        if method == "POST" and path == "run":
+                            raise AmbiguousSubmission("Runpod submission did not return a job record; check the provider before retrying.")
+                        raise RunpodError("Runpod returned an invalid response")
+                    return result
+            except (asyncio.TimeoutError, aiohttp.ClientError, OSError):
+                if method == "POST" and path == "run":
+                    raise AmbiguousSubmission("Runpod submission connection was interrupted; acceptance is unknown. Check the provider before retrying.") from None
+                if method == "GET" and attempt + 1 < attempts:
+                    await asyncio.sleep(0.5 * (2 ** attempt))
+                    continue
+                raise RunpodError(f"Runpod {path.split('/')[0]} connection failed") from None
+        raise RunpodError("Runpod request failed")
+
+    async def health(self):
+        return await self._request("GET", "health")
+
+    async def submit(self, workflow, images, timeout_seconds):
+        body = serialize_request(workflow, images, timeout_seconds)
+        result = await self._request("POST", "run", request_body=body)
+        try:
+            return checked_id(result.get("id"))
+        except ValueError:
+            raise AmbiguousSubmission("Runpod submission returned no valid job ID; check the provider before retrying.") from None
+
+    async def status(self, job_id):
+        return await self._request("GET", f"status/{checked_id(job_id)}")
+
+    async def cancel(self, job_id):
+        return await self._request("POST", f"cancel/{checked_id(job_id)}")

api/orchestration/serverless_ledger.py

@@ -0,0 +1,57 @@
+"""Private, atomic job-ID journal. Contains identifiers/status only, never payloads."""
+import json
+import os
+from pathlib import Path
+import tempfile
+import threading
+
+from .serverless_state import TERMINAL
+FIELDS = {"record_id", "execution_id", "worker_id", "participant_id", "endpoint_id", "job_id", "status", "created_at", "updated_at", "error", "prompt_id"}
+
+
+class JobLedger:
+    def __init__(self, path):
+        self.path = Path(path)
+        self._lock = threading.RLock()
+
+    def _read(self):
+        if self.path.is_symlink():
+            raise RuntimeError("Unsafe serverless job journal path")
+        if not self.path.exists():
+            return {}
+        try:
+            data = json.loads(self.path.read_text(encoding="utf-8"))
+            if not isinstance(data, dict) or any(not isinstance(v, dict) or set(v) - FIELDS for v in data.values()):
+                raise ValueError()
+            return data
+        except (ValueError, OSError):
+            raise RuntimeError("Serverless job journal is unreadable; preserve it for recovery before submitting more work") from None
+
+    def all(self):
+        with self._lock:
+            return list(self._read().values())
+
+    def update(self, record_id, **fields):
+        if set(fields) - FIELDS:
+            raise ValueError("Unexpected serverless journal field")
+        with self._lock:
+            data = self._read()
+            data.setdefault(record_id, {"record_id": record_id}).update(fields)
+            # Keep every unresolved job, plus a bounded terminal history.
+            completed = sorted((k for k, v in data.items() if v.get("status") in TERMINAL), key=lambda k: data[k].get("updated_at", 0))
+            for key in completed[:-200]:
+                data.pop(key)
+            if self.path.parent.is_symlink():
+                raise RuntimeError("Unsafe serverless journal directory")
+            self.path.parent.mkdir(mode=0o700, parents=True, exist_ok=True)
+            fd, temporary = tempfile.mkstemp(prefix=".jobs-", dir=self.path.parent)
+            try:
+                with os.fdopen(fd, "w", encoding="utf-8") as stream:
+                    json.dump(data, stream, allow_nan=False)
+                    stream.flush()
+                    os.fsync(stream.fileno())
+                os.replace(temporary, self.path)
+            finally:
+                if os.path.exists(temporary):
+                    os.unlink(temporary)
+            return dict(data[record_id])

api/orchestration/serverless_media.py

@@ -0,0 +1,601 @@
+"""Bounded image transport for generic Runpod-backed ComfyUI jobs."""
+
+import asyncio
+import base64
+import binascii
+import hashlib
+import io
+import ipaddress
+import os
+import re
+import socket
+import stat
+import warnings
+from pathlib import Path, PurePosixPath
+from urllib.parse import urljoin, urlsplit
+
+import aiohttp
+import numpy as np
+from PIL import Image, UnidentifiedImageError
+
+
+MAX_INPUT_IMAGE_BYTES = 25 * 1024 * 1024
+MAX_INPUT_TOTAL_BYTES = 10_000_000
+MAX_OUTPUT_FILES = 1024
+MAX_OUTPUT_IMAGE_BYTES = 128 * 1024 * 1024
+MAX_OUTPUT_TOTAL_BYTES = 256 * 1024 * 1024
+MAX_IMAGE_PIXELS = 40_000_000
+MAX_OUTPUT_ARRAY_BYTES = 256 * 1024 * 1024
+MAX_REDIRECTS = 3
+_DOWNLOAD_TIMEOUT_SECONDS = 60
+_CHUNK_SIZE = 64 * 1024
+
+_IMAGE_EXTENSIONS = {".png", ".jpg", ".jpeg", ".webp", ".bmp"}
+_UNSUPPORTED_IMAGE_EXTENSIONS = {".gif", ".tif", ".tiff", ".avif", ".heic", ".ico", ".svg"}
+_AUDIO_EXTENSIONS = {".wav", ".mp3", ".flac", ".m4a", ".aac", ".ogg", ".opus", ".aiff", ".aif", ".wma"}
+_VIDEO_EXTENSIONS = {".mp4", ".avi", ".mov", ".mkv", ".webm"}
+_MEDIA_EXTENSIONS = _IMAGE_EXTENSIONS | _UNSUPPORTED_IMAGE_EXTENSIONS | _AUDIO_EXTENSIONS | _VIDEO_EXTENSIONS
+_MEDIA_ANNOTATION = re.compile(r"\s+\[([A-Za-z0-9_-]{1,64})\]$")
+_PREFIX_RE = re.compile(r"[A-Za-z0-9][A-Za-z0-9_.-]{0,199}\Z", re.ASCII)
+_SAFE_FILENAME_RE = re.compile(r"[A-Za-z0-9_.-]{1,255}\Z", re.ASCII)
+
+
+class ServerlessMediaError(ValueError):
+    """Sanitized input, artifact, or download validation failure."""
+
+
+def _folder_roots(folder_paths):
+    roots = {}
+    for kind, getter_name in (
+        ("input", "get_input_directory"),
+        ("output", "get_output_directory"),
+        ("temp", "get_temp_directory"),
+    ):
+        getter = getattr(folder_paths, getter_name, None)
+        if not callable(getter):
+            raise ServerlessMediaError("ComfyUI media directories are unavailable")
+        try:
+            roots[kind] = Path(getter()).resolve(strict=True)
+        except Exception:
+            raise ServerlessMediaError("ComfyUI media directories are unavailable") from None
+    return roots
+
+
+def _inside_roots(path, roots):
+    for root in roots.values():
+        try:
+            path.relative_to(root)
+            return True
+        except ValueError:
+            continue
+    return False
+
+
+def _normalize_media_reference(value):
+    if not isinstance(value, str) or not value or value != value.strip() or "\x00" in value:
+        raise ServerlessMediaError("Image input reference is invalid")
+    value = value.replace("\\", "/")
+    explicit_root = None
+    annotation = _MEDIA_ANNOTATION.search(value)
+    if annotation:
+        marker = annotation.group(1).lower()
+        if marker not in {"input", "output", "temp"}:
+            raise ServerlessMediaError("Image input annotation is unsupported")
+        explicit_root = marker
+        value = value[: annotation.start()]
+    if "://" in value:
+        raise ServerlessMediaError("Remote media references are not supported as image inputs")
+    if re.match(r"^[A-Za-z]:", value) or value.startswith("//"):
+        raise ServerlessMediaError("Image input path is outside the allowed ComfyUI media directories")
+
+    path = PurePosixPath(value)
+    if any(part == ".." for part in path.parts):
+        raise ServerlessMediaError("Image input path traversal is not allowed")
+    if not path.parts or value in {".", ""}:
+        raise ServerlessMediaError("Image input reference is invalid")
+
+    if explicit_root:
+        value = f"{value} [{explicit_root}]"
+    suffix = Path(path.name).suffix.lower()
+    return value, suffix
+
+
+def _resolve_input_path(value, suffix, folder_paths, roots):
+    if suffix not in _IMAGE_EXTENSIONS:
+        if suffix in _MEDIA_EXTENSIONS:
+            raise ServerlessMediaError("Audio and video inputs are not supported by serverless image transport")
+        raise ServerlessMediaError("Unsupported serverless image file type")
+
+    try:
+        source = Path(value)
+        if source.is_absolute():
+            annotation = _MEDIA_ANNOTATION.search(value)
+            unannotated = value[:annotation.start()] if annotation else value
+            resolved = Path(unannotated).resolve(strict=True)
+        else:
+            resolver = getattr(folder_paths, "get_annotated_filepath", None)
+            if not callable(resolver):
+                raise ServerlessMediaError("ComfyUI media path resolver is unavailable")
+            resolved = Path(resolver(value)).resolve(strict=True)
+    except ServerlessMediaError:
+        raise
+    except Exception:
+        raise ServerlessMediaError("Image input file is missing or invalid") from None
+
+    if not _inside_roots(resolved, roots):
+        raise ServerlessMediaError("Image input path is outside the allowed ComfyUI media directories")
+    try:
+        mode = resolved.stat().st_mode
+    except OSError:
+        raise ServerlessMediaError("Image input file is missing or invalid") from None
+    if not stat.S_ISREG(mode):
+        raise ServerlessMediaError("Image input is not a regular file")
+    return resolved
+
+
+def _iter_image_references(workflow):
+    if not isinstance(workflow, dict):
+        raise ServerlessMediaError("Serverless workflow must be an API graph")
+    references = []
+    for node_id, node in workflow.items():
+        if not isinstance(node, dict):
+            continue
+        class_type = node.get("class_type")
+        inputs = node.get("inputs", {})
+        if class_type == "LoadImage" and not isinstance(inputs, dict):
+            raise ServerlessMediaError("LoadImage requires a literal image filename; linked image inputs are unsupported")
+        if not isinstance(inputs, dict):
+            continue
+        if class_type == "LoadImage":
+            value = inputs.get("image")
+            if not isinstance(value, str):
+                raise ServerlessMediaError("LoadImage requires a literal image filename; linked image inputs are unsupported")
+            references.append((node_id, inputs, "image", value, True))
+
+        for key in ("image", "file"):
+            if key == "image" and class_type == "LoadImage":
+                continue
+            value = inputs.get(key)
+            if not isinstance(value, str):
+                continue
+            normalized, suffix = _normalize_media_reference(value)
+            if suffix in _IMAGE_EXTENSIONS:
+                references.append((node_id, inputs, key, normalized, False))
+            elif suffix in _AUDIO_EXTENSIONS | _VIDEO_EXTENSIONS:
+                raise ServerlessMediaError("Audio and video inputs are not supported by serverless image transport")
+            elif suffix in _UNSUPPORTED_IMAGE_EXTENSIONS or (key == "image" and suffix):
+                raise ServerlessMediaError("Unsupported serverless image file type")
+
+        for key in ("audio", "video"):
+            value = inputs.get(key)
+            if not isinstance(value, str):
+                continue
+            _normalized, suffix = _normalize_media_reference(value)
+            if suffix in _MEDIA_EXTENSIONS:
+                raise ServerlessMediaError("Audio and video inputs are not supported by serverless image transport")
+    return references
+
+
+def _read_limited(path, expected_size):
+    try:
+        descriptor = os.open(path, os.O_RDONLY | getattr(os, "O_NOFOLLOW", 0))
+        with os.fdopen(descriptor, "rb") as handle:
+            info = os.fstat(handle.fileno())
+            if not stat.S_ISREG(info.st_mode) or info.st_size > MAX_INPUT_IMAGE_BYTES:
+                raise ServerlessMediaError("Image input exceeds the permitted size")
+            if info.st_size > MAX_INPUT_TOTAL_BYTES:
+                raise ServerlessMediaError(
+                    f"Total raw serverless image input size exceeds {MAX_INPUT_TOTAL_BYTES:,} bytes; "
+                    "base64 encoding adds about 33%, and the complete request body is also limited to 10,000,000 bytes"
+                )
+            data = handle.read(MAX_INPUT_TOTAL_BYTES + 1)
+    except ServerlessMediaError:
+        raise
+    except OSError:
+        raise ServerlessMediaError("Image input file could not be read safely") from None
+    if len(data) > MAX_INPUT_TOTAL_BYTES:
+        raise ServerlessMediaError(
+            f"Total raw serverless image input size exceeds {MAX_INPUT_TOTAL_BYTES:,} bytes; "
+            "base64 encoding adds about 33%, and the complete request body is also limited to 10,000,000 bytes"
+        )
+    if len(data) != expected_size:
+        raise ServerlessMediaError("Image input changed while being read")
+    return data
+
+
+async def prepare_images(workflow: dict) -> list[dict]:
+    """Pack referenced local images and rewrite only their graph input references.
+
+    ``workflow`` is expected to be a private API-graph copy owned by the caller.
+    Returned ``name`` values are flat generated names suitable for ``input.images``.
+    """
+    references = _iter_image_references(workflow)
+    if not references:
+        return []
+
+    # Import only in ComfyUI, keeping this module importable in the serverless client.
+    try:
+        import folder_paths
+    except Exception:
+        raise ServerlessMediaError("ComfyUI media directories are unavailable") from None
+
+    roots = _folder_roots(folder_paths)
+    resolved_refs = []
+    unique_paths = {}
+    total_source_bytes = 0
+    for node_id, inputs, key, value, _is_load_image in references:
+        normalized, suffix = _normalize_media_reference(value)
+        if key == "image" and suffix not in _IMAGE_EXTENSIONS:
+            if suffix in _AUDIO_EXTENSIONS | _VIDEO_EXTENSIONS:
+                raise ServerlessMediaError("Audio and video inputs are not supported by serverless image transport")
+            raise ServerlessMediaError("Unsupported serverless image file type")
+        path = _resolve_input_path(normalized, suffix, folder_paths, roots)
+        if path not in unique_paths:
+            try:
+                size = path.stat().st_size
+            except OSError:
+                raise ServerlessMediaError("Image input file is missing or invalid") from None
+            if size < 0 or size > MAX_INPUT_IMAGE_BYTES:
+                raise ServerlessMediaError("Image input exceeds the permitted size")
+            total_source_bytes += size
+            if total_source_bytes > MAX_INPUT_TOTAL_BYTES:
+                raise ServerlessMediaError(
+                    f"Total raw serverless image input size is {total_source_bytes:,} bytes; "
+                    f"the input read guard is {MAX_INPUT_TOTAL_BYTES:,} bytes. Base64 encoding "
+                    "adds about 33%, and the complete request body is also limited to 10,000,000 bytes"
+                )
+            unique_paths[path] = {"size": size, "suffix": suffix}
+        resolved_refs.append((inputs, key, path))
+
+    packed_by_name = {}
+    path_to_name = {}
+    for path, metadata in unique_paths.items():
+        data = _read_limited(path, metadata["size"])
+        digest = hashlib.sha256(data).hexdigest()
+        name = f"serverless_{digest}{metadata['suffix']}"
+        path_to_name[path] = name
+        packed_by_name.setdefault(name, {"name": name, "image": base64.b64encode(data).decode("ascii")})
+
+    for inputs, key, path in resolved_refs:
+        inputs[key] = path_to_name[path]
+    return [packed_by_name[name] for name in sorted(packed_by_name)]
+
+
+def _validate_artifact_url(url):
+    if not isinstance(url, str) or not url or any(ord(char) < 0x20 or char.isspace() for char in url):
+        raise ServerlessMediaError("Serverless image artifact URL is invalid")
+    try:
+        parsed = urlsplit(url)
+        port = parsed.port
+    except ValueError:
+        raise ServerlessMediaError("Serverless image artifact URL is invalid") from None
+    if (
+        parsed.scheme.lower() != "https"
+        or not parsed.hostname
+        or parsed.username is not None
+        or parsed.password is not None
+        or port not in (None, 443)
+        or parsed.fragment
+    ):
+        raise ServerlessMediaError("Serverless image artifact URL must be a public HTTPS URL")
+    host = parsed.hostname.rstrip(".").lower()
+    if host in {"localhost", "metadata", "metadata.google.internal"} or host.endswith((".localhost", ".local", ".internal")):
+        raise ServerlessMediaError("Serverless image artifact URL must be a public HTTPS URL")
+    try:
+        address = ipaddress.ip_address(host)
+    except ValueError:
+        return url
+    if not address.is_global:
+        raise ServerlessMediaError("Serverless image artifact URL must resolve to a public address")
+    return url
+
+
+class _PublicResolver(aiohttp.abc.AbstractResolver):
+    """Resolve once, reject non-public answers, and pin the checked address."""
+
+    async def _getaddrinfo(self, host, port, family):
+        return await asyncio.get_running_loop().getaddrinfo(
+            host, port, family=family, type=socket.SOCK_STREAM
+        )
+
+    async def resolve(self, host, port=0, family=socket.AF_INET):
+        try:
+            results = await self._getaddrinfo(host, port, family)
+        except (OSError, asyncio.TimeoutError):
+            raise OSError("Public artifact host could not be resolved") from None
+        if not results:
+            raise OSError("Public artifact host could not be resolved")
+
+        resolved = []
+        seen = set()
+        for result_family, _socktype, proto, _canonname, sockaddr in results:
+            try:
+                address = ipaddress.ip_address(sockaddr[0])
+            except (ValueError, TypeError, IndexError):
+                raise OSError("Public artifact host resolved to an invalid address") from None
+            if not address.is_global:
+                raise OSError("Public artifact host resolved to a non-public address")
+            key = (str(address), result_family)
+            if key in seen:
+                continue
+            seen.add(key)
+            resolved.append({
+                "hostname": host,
+                "host": str(address),
+                "port": port,
+                "family": result_family,
+                "proto": proto,
+                "flags": 0,
+            })
+        if not resolved:
+            raise OSError("Public artifact host could not be resolved")
+        return resolved
+
+    async def close(self):
+        return None
+
+
+async def download_public_artifact(url: str, max_bytes: int = MAX_OUTPUT_IMAGE_BYTES) -> bytes:
+    """Download a signed public HTTPS artifact without provider credentials."""
+    if isinstance(max_bytes, bool) or not isinstance(max_bytes, int) or not 0 < max_bytes <= MAX_OUTPUT_IMAGE_BYTES:
+        raise ServerlessMediaError("Serverless image artifact exceeds the permitted size")
+    current_url = _validate_artifact_url(url)
+    resolver = _PublicResolver()
+    connector = aiohttp.TCPConnector(resolver=resolver, use_dns_cache=False, limit=4)
+    timeout = aiohttp.ClientTimeout(
+        total=_DOWNLOAD_TIMEOUT_SECONDS,
+        connect=10,
+        sock_read=15,
+    )
+    try:
+        async with aiohttp.ClientSession(
+            connector=connector,
+            timeout=timeout,
+            trust_env=False,
+            cookie_jar=aiohttp.DummyCookieJar(),
+        ) as session:
+            for redirect_count in range(MAX_REDIRECTS + 1):
+                _validate_artifact_url(current_url)
+                async with session.get(current_url, allow_redirects=False) as response:
+                    if response.status in {301, 302, 303, 307, 308}:
+                        location = response.headers.get("Location")
+                        if not location or redirect_count >= MAX_REDIRECTS:
+                            raise ServerlessMediaError("Serverless image artifact redirect is invalid")
+                        current_url = _validate_artifact_url(urljoin(current_url, location))
+                        continue
+                    if response.status != 200:
+                        raise ServerlessMediaError("Serverless image artifact download failed")
+                    payload = bytearray()
+                    async for chunk in response.content.iter_chunked(_CHUNK_SIZE):
+                        if not isinstance(chunk, bytes):
+                            raise ServerlessMediaError("Serverless image artifact response is invalid")
+                        if len(payload) + len(chunk) > max_bytes:
+                            raise ServerlessMediaError("Serverless image artifact exceeds the permitted size")
+                        payload.extend(chunk)
+                    if not payload:
+                        raise ServerlessMediaError("Serverless image artifact is empty")
+                    return bytes(payload)
+    except ServerlessMediaError:
+        raise
+    except Exception:
+        raise ServerlessMediaError("Serverless image artifact download failed") from None
+    raise ServerlessMediaError("Serverless image artifact download failed")
+
+
+def _validated_manifest(manifest):
+    if not isinstance(manifest, list) or len(manifest) > MAX_OUTPUT_FILES:
+        raise ServerlessMediaError("Serverless image manifest is invalid")
+    normalized = []
+    seen_prefixes = set()
+    prefixes = []
+    for item in manifest:
+        if not isinstance(item, dict):
+            raise ServerlessMediaError("Serverless image manifest is invalid")
+        if item.get("kind") != "image":
+            raise ServerlessMediaError("Only image collectors are supported by serverless media transport")
+        for key in ("collector_id", "job_id", "node_id"):
+            value = item.get(key)
+            if not isinstance(value, (str, int)) or isinstance(value, bool) or not str(value):
+                raise ServerlessMediaError("Serverless image manifest is invalid")
+        prefix = item.get("prefix")
+        if not isinstance(prefix, str) or not _PREFIX_RE.fullmatch(prefix) or prefix in {".", ".."}:
+            raise ServerlessMediaError("Serverless image manifest contains an unsafe filename prefix")
+        folded = prefix.casefold()
+        if folded in seen_prefixes:
+            raise ServerlessMediaError("Serverless image manifest prefixes must be unique")
+        seen_prefixes.add(folded)
+        prefixes.append(prefix)
+        normalized.append({
+            "collector_id": str(item["collector_id"]),
+            "job_id": str(item["job_id"]),
+            "node_id": str(item["node_id"]),
+            "prefix": prefix,
+        })
+    for first in prefixes:
+        for second in prefixes:
+            if first != second and second.casefold().startswith(first.casefold() + "_"):
+                raise ServerlessMediaError("Serverless image manifest prefixes are ambiguous")
+    return normalized
+
+
+def _artifact_filename_matches(filename, prefixes):
+    if not isinstance(filename, str):
+        return None
+    basename = filename.replace("\\", "/").rsplit("/", 1)[-1]
+    matches = [item for item in prefixes if basename.startswith(item["prefix"] + "_")]
+    if len(matches) > 1:
+        raise ServerlessMediaError("Serverless image artifact matches multiple collectors")
+    return matches[0] if matches else None
+
+
+def _canonical_artifact_identity(entry, filename):
+    artifact_path = entry.get("artifact_path")
+    artifact_type = entry.get("artifact_type", "output")
+    if artifact_path is not None:
+        if not isinstance(artifact_path, str) or not artifact_path or "\x00" in artifact_path:
+            raise ServerlessMediaError("Serverless image artifact identity is invalid")
+        normalized_path = os.path.normcase(os.path.normpath(artifact_path.replace("\\", "/")))
+    else:
+        subfolder = entry.get("subfolder", "")
+        if not isinstance(subfolder, str) or "\x00" in subfolder:
+            raise ServerlessMediaError("Serverless image artifact identity is invalid")
+        normalized_path = os.path.normcase(os.path.normpath(f"{subfolder}/{filename}"))
+    return str(artifact_type), normalized_path
+
+
+def _decode_base64_data(value, max_bytes=MAX_OUTPUT_IMAGE_BYTES):
+    if not isinstance(value, str) or not value:
+        raise ServerlessMediaError("Serverless image artifact data is invalid")
+    max_bytes = min(max_bytes, MAX_OUTPUT_IMAGE_BYTES)
+    max_encoded = ((max_bytes + 2) // 3) * 4
+    if len(value) > max_encoded:
+        raise ServerlessMediaError("Serverless image artifact exceeds the permitted size")
+    padding = len(value) - len(value.rstrip("="))
+    estimated_size = (len(value) // 4) * 3 - padding
+    if estimated_size > max_bytes:
+        raise ServerlessMediaError("Serverless image artifact exceeds the permitted size")
+    try:
+        payload = base64.b64decode(value, validate=True)
+    except (binascii.Error, ValueError):
+        raise ServerlessMediaError("Serverless image artifact data is invalid") from None
+    if not payload:
+        raise ServerlessMediaError("Serverless image artifact is empty")
+    if len(payload) > max_bytes:
+        raise ServerlessMediaError("Serverless image artifact exceeds the permitted size")
+    return payload
+
+
+def _decode_image_array(payload, total_array_bytes):
+    try:
+        with warnings.catch_warnings():
+            warnings.simplefilter("error", Image.DecompressionBombWarning)
+            image = Image.open(io.BytesIO(payload))
+            width, height = image.size
+            if width <= 0 or height <= 0 or width * height > MAX_IMAGE_PIXELS:
+                raise ServerlessMediaError("Serverless image dimensions exceed the permitted limit")
+            has_alpha = image.mode in {"RGBA", "LA"} or "transparency" in image.info
+            mode = "RGBA" if has_alpha else "RGB"
+            output_bytes = width * height * (4 if has_alpha else 3) * 4
+            if total_array_bytes + output_bytes > MAX_OUTPUT_ARRAY_BYTES:
+                raise ServerlessMediaError("Decoded serverless images exceed the permitted memory limit")
+            converted = image.convert(mode)
+            converted.load()
+            array = np.array(converted, dtype=np.float32, copy=True)
+            np.multiply(array, 1.0 / 255.0, out=array)
+            return array, output_bytes
+    except ServerlessMediaError:
+        raise
+    except (UnidentifiedImageError, OSError, ValueError, Image.DecompressionBombError, Image.DecompressionBombWarning):
+        raise ServerlessMediaError("Serverless image artifact could not be decoded") from None
+    except Exception:
+        raise ServerlessMediaError("Serverless image artifact could not be decoded") from None
+
+
+def _image_output(output):
+    """Unwrap handler output and reject explicit execution errors."""
+    if not isinstance(output, dict):
+        raise ServerlessMediaError("Serverless job output is invalid")
+    nested = output.get("output")
+    data = nested if isinstance(nested, dict) else output
+    for candidate in (output, data):
+        if candidate.get("error") or candidate.get("errors"):
+            raise ServerlessMediaError("Serverless job reported an execution error")
+
+    return data
+
+
+async def _select_image_payloads(data, expected, download):
+    """Match artifact identities and fetch bounded bytes, without image decoding."""
+    artifacts = data.get("images", [])
+    if not isinstance(artifacts, list) or len(artifacts) > MAX_OUTPUT_FILES:
+        raise ServerlessMediaError("Serverless job returned too many image files")
+
+    selected = {item["prefix"]: [] for item in expected}
+    identities = set()
+    total_payload_bytes = 0
+    for entry in artifacts:
+        if not isinstance(entry, dict):
+            continue
+        filename = entry.get("filename")
+        manifest_entry = _artifact_filename_matches(filename, expected)
+        if manifest_entry is None:
+            continue
+        normalized_filename = filename.replace("\\", "/")
+        if "/" in normalized_filename or not _SAFE_FILENAME_RE.fullmatch(filename):
+            raise ServerlessMediaError("Serverless image artifact filename is invalid")
+        if Path(filename).suffix.lower() not in _IMAGE_EXTENSIONS:
+            raise ServerlessMediaError("Serverless image artifact has an unsupported file type")
+        if "artifact_type" in entry and entry["artifact_type"] != "output":
+            raise ServerlessMediaError("Expected serverless image artifact was not an output image")
+        identity = _canonical_artifact_identity(entry, filename)
+        if identity in identities:
+            raise ServerlessMediaError("Duplicate serverless image artifact identity")
+        identities.add(identity)
+        match = re.match(rf"^{re.escape(manifest_entry['prefix'])}_(\d+)(?:_|\.)", filename)
+        if match is None:
+            raise ServerlessMediaError("Expected serverless image filename counter is invalid")
+        artifact_type = entry.get("type")
+        remaining_bytes = MAX_OUTPUT_TOTAL_BYTES - total_payload_bytes
+        if remaining_bytes <= 0:
+            raise ServerlessMediaError("Serverless image artifacts exceed the permitted total size")
+        if artifact_type == "base64":
+            payload = _decode_base64_data(entry.get("data"), remaining_bytes)
+        elif artifact_type in {"gcs_url", "s3_url", "url"}:
+            url = entry.get("data") or entry.get("url")
+            url = _validate_artifact_url(url)
+            try:
+                if download is None:
+                    payload = await download_public_artifact(
+                        url, max_bytes=min(MAX_OUTPUT_IMAGE_BYTES, remaining_bytes)
+                    )
+                else:
+                    payload = await download(url)
+            except ServerlessMediaError:
+                raise
+            except Exception:
+                raise ServerlessMediaError("Serverless image artifact download failed") from None
+            if not isinstance(payload, bytes) or not payload:
+                raise ServerlessMediaError("Serverless image artifact response is invalid")
+            if len(payload) > MAX_OUTPUT_IMAGE_BYTES:
+                raise ServerlessMediaError("Serverless image artifact exceeds the permitted size")
+        else:
+            raise ServerlessMediaError("Serverless image artifact transport type is unsupported")
+        total_payload_bytes += len(payload)
+        if total_payload_bytes > MAX_OUTPUT_TOTAL_BYTES:
+            raise ServerlessMediaError("Serverless image artifacts exceed the permitted total size")
+        selected[manifest_entry["prefix"]].append((int(match.group(1)), filename, identity, payload))
+
+    if any(not selected[item["prefix"]] for item in expected):
+        raise ServerlessMediaError("Serverless job is missing an expected collector image")
+    return selected
+
+
+def _decode_image_groups(selected, expected):
+    """Decode in stable counter order, bounding array memory across collectors."""
+    result = {}
+    total_array_bytes = 0
+    for manifest_entry in expected:
+        prefix = manifest_entry["prefix"]
+        rows = sorted(selected[prefix], key=lambda row: (row[0], row[1], row[2]))
+        arrays = []
+        shape = None
+        for _counter, _filename, _identity, payload in rows:
+            array, added_bytes = _decode_image_array(payload, total_array_bytes)
+            total_array_bytes += added_bytes
+            if shape is None:
+                shape = array.shape
+            elif array.shape != shape:
+                raise ServerlessMediaError("Images from one serverless collector have inconsistent dimensions")
+            arrays.append(array)
+        result.setdefault(manifest_entry["job_id"], []).extend(arrays)
+    return result
+
+
+async def collect_images(output: dict, manifest: list, download=None) -> dict[str, list[np.ndarray]]:
+    """Validate and decode expected image artifacts, grouped by serverless job ID."""
+    data = _image_output(output)
+    expected = _validated_manifest(manifest)
+    if not expected:
+        return {}
+    selected = await _select_image_payloads(data, expected, download)
+    return _decode_image_groups(selected, expected)

api/orchestration/serverless_prompt.py

@@ -0,0 +1,420 @@
+"""Compile prepared distributed prompts into worker-side ComfyUI image jobs.
+
+This module is deliberately pure: it only copies and rewrites an API prompt
+mapping and never imports or executes ComfyUI nodes.
+"""
+
+from __future__ import annotations
+
+import copy
+import hashlib
+import json
+import re
+
+
+_SUPPORTED_DISTRIBUTED = {
+    "DistributedCollector",
+    "DistributedSeed",
+    "DistributedValue",
+    "DistributedModelName",
+}
+_REMOVED_SCALAR_NODES = {
+    "DistributedSeed",
+    "DistributedValue",
+    "DistributedModelName",
+}
+_PRIMITIVE_NODES = {
+    "PrimitiveBoolean": False,
+    "PrimitiveFloat": 0.0,
+    "PrimitiveInt": 0,
+    "PrimitiveString": "",
+}
+_REMOVED_NODE_OUTPUTS = {
+    "DistributedCollector": {0, 1},
+    "DistributedSeed": {0},
+    "DistributedValue": {0},
+    "DistributedModelName": {0},
+}
+_SEED_DEFAULT = 1125899906842
+
+
+def _is_link(value):
+    return isinstance(value, list) and len(value) == 2
+
+
+def _is_unsupported_distributed(class_type):
+    if not isinstance(class_type, str) or class_type in _SUPPORTED_DISTRIBUTED:
+        return False
+    lowered = class_type.lower()
+    return "distributed" in lowered or "batchdivider" in lowered
+
+
+def _safe_component(value):
+    cleaned = re.sub(r"[^A-Za-z0-9_-]+", "_", str(value)).strip("_-")
+    return cleaned[:24] or "empty"
+
+
+def _filename_prefix(job_id, worker_id, collector_id, kind):
+    identity = json.dumps(
+        [job_id, worker_id, collector_id, kind],
+        ensure_ascii=False,
+        separators=(",", ":"),
+    )
+    digest = hashlib.sha256(identity.encode("utf-8")).hexdigest()
+    parts = (
+        "serverless",
+        _safe_component(job_id),
+        _safe_component(worker_id),
+        _safe_component(collector_id),
+        kind,
+        digest,
+    )
+    return "_".join(parts)
+
+
+def _worker_index(worker_id):
+    """Mirror nodes.utilities worker-index parsing; None means fallback."""
+    try:
+        if worker_id.startswith("worker_"):
+            return int(worker_id.split("_")[1])
+        return int(worker_id)
+    except (ValueError, IndexError):
+        return None
+
+
+def _coerce(value, value_type):
+    if value_type == "INT":
+        return int(float(value))
+    if value_type == "FLOAT":
+        return float(value)
+    return value
+
+
+def _coerce_safe(value, value_type):
+    try:
+        return _coerce(value, value_type)
+    except (TypeError, ValueError):
+        return value
+
+
+def _stringify_model_name(value):
+    if isinstance(value, str):
+        return value
+    if isinstance(value, (int, float, bool)):
+        return str(value)
+    try:
+        return json.dumps(value, indent=4)
+    except (TypeError, ValueError):
+        return str(value)
+
+
+def _model_name_result(text):
+    if isinstance(text, list):
+        values = [_stringify_model_name(value) for value in text]
+        return values[0] if len(values) == 1 else values
+    return _stringify_model_name(text)
+
+
+def _validate_prepared_prompt(prepared_prompt):
+    """Validate graph shape, links, collectors, and supported topology."""
+    if not isinstance(prepared_prompt, dict):
+        raise ValueError("prepared prompt must be an API graph mapping")
+
+    source_nodes = {}
+    for raw_id, node in prepared_prompt.items():
+        node_id = str(raw_id)
+        if node_id in source_nodes:
+            raise ValueError("prompt has duplicate node IDs after string normalization")
+        if not isinstance(node, dict):
+            raise ValueError(f"prompt node {node_id!r} must be an object")
+        class_type = node.get("class_type")
+        if not isinstance(class_type, str):
+            raise ValueError(f"prompt node {node_id!r} is missing a class_type")
+        if _is_unsupported_distributed(class_type):
+            raise ValueError(f"unsupported distributed or batch node: {class_type}")
+        inputs = node.get("inputs", {})
+        if not isinstance(inputs, dict):
+            raise ValueError(f"prompt node {node_id!r} inputs must be an object")
+        source_nodes[node_id] = node
+
+    if not source_nodes:
+        raise ValueError("prepared prompt contains no nodes")
+
+    # Validate every API link before any graph rewriting. Slots on the nodes
+    # removed by this compiler have a known schema and are checked exactly.
+    for node_id, node in source_nodes.items():
+        for input_name, value in node.get("inputs", {}).items():
+            if not _is_link(value):
+                continue
+            upstream_id, output_slot = value
+            upstream_id = str(upstream_id)
+            if upstream_id not in source_nodes:
+                raise ValueError(
+                    f"invalid link at node {node_id!r} input {input_name!r}: "
+                    f"source {upstream_id!r} does not exist"
+                )
+            if isinstance(output_slot, bool) or not isinstance(output_slot, int) or output_slot < 0:
+                raise ValueError(
+                    f"invalid link at node {node_id!r} input {input_name!r}: "
+                    "output slot must be a nonnegative integer"
+                )
+            upstream_class = source_nodes[upstream_id]["class_type"]
+            valid_slots = _REMOVED_NODE_OUTPUTS.get(upstream_class)
+            if valid_slots is not None and output_slot not in valid_slots:
+                raise ValueError(
+                    f"invalid output slot {output_slot} for removed node "
+                    f"{upstream_class} {upstream_id!r}"
+                )
+
+    collector_ids = sorted(
+        (node_id for node_id, node in source_nodes.items() if node["class_type"] == "DistributedCollector"),
+        key=str,
+    )
+    if not collector_ids:
+        raise ValueError("prepared prompt contains no DistributedCollector nodes")
+
+    # Only image collectors are supported by the current generic handler.
+    for collector_id in collector_ids:
+        inputs = source_nodes[collector_id].get("inputs", {})
+        if "audio" in inputs:
+            raise ValueError("collector audio input is unsupported by this image-only compiler")
+        if not _is_link(inputs.get("images")):
+            raise ValueError(f"collector {collector_id!r} is missing a linked image input")
+        source_id = str(inputs["images"][0])
+        source_class = source_nodes[source_id]["class_type"]
+        if source_class in _REMOVED_SCALAR_NODES:
+            raise ValueError(f"collector {collector_id!r} image input must link to an image node")
+
+    # A collector cannot be used, directly or through processing, as another
+    # collector's input: this first compiler handles disconnected jobs only.
+    for collector_id in collector_ids:
+        pending = [str(source_nodes[collector_id]["inputs"]["images"][0])]
+        seen = set()
+        while pending:
+            node_id = pending.pop()
+            if node_id in seen:
+                continue
+            seen.add(node_id)
+            if source_nodes[node_id]["class_type"] == "DistributedCollector":
+                raise ValueError("collector-to-collector dependence is unsupported")
+            for value in source_nodes[node_id].get("inputs", {}).values():
+                if _is_link(value):
+                    pending.append(str(value[0]))
+
+    return source_nodes, collector_ids
+
+
+def _make_scalar_resolver(source_nodes):
+    """Build pure resolvers for distributed scalar nodes and stock primitives."""
+    # Resolve only stock primitive nodes. A visited set makes malformed
+    # primitive-to-primitive cycles fail clearly instead of recursing forever.
+    def resolve_primitive(node_id, visiting=None):
+        node_id = str(node_id)
+        node = source_nodes[node_id]
+        class_type = node["class_type"]
+        if class_type not in _PRIMITIVE_NODES:
+            raise ValueError("linked value must be literal or supported PrimitiveInt/Float/String/Boolean")
+        if visiting is None:
+            visiting = set()
+        if node_id in visiting:
+            raise ValueError(f"cycle while resolving primitive value at node {node_id!r}")
+        visiting.add(node_id)
+        value = node.get("inputs", {}).get("value", _PRIMITIVE_NODES[class_type])
+        if _is_link(value):
+            parent_id, output_slot = str(value[0]), value[1]
+            if output_slot != 0:
+                raise ValueError(f"invalid output slot {output_slot} for primitive node {parent_id!r}")
+            if source_nodes[parent_id]["class_type"] not in _PRIMITIVE_NODES:
+                raise ValueError("linked primitive value must come from a stock Primitive node")
+            result = resolve_primitive(parent_id, visiting)
+        else:
+            result = copy.deepcopy(value)
+        visiting.remove(node_id)
+        return result
+
+    def resolve_literal_or_primitive(value, context):
+        if _is_link(value):
+            source_id, output_slot = str(value[0]), value[1]
+            if output_slot != 0 or source_nodes[source_id]["class_type"] not in _PRIMITIVE_NODES:
+                raise ValueError(f"{context} must be a literal or supported Primitive link")
+            try:
+                return resolve_primitive(source_id)
+            except ValueError as exc:
+                if "cycle" in str(exc):
+                    raise
+                raise ValueError(f"{context} must be a literal or supported Primitive link") from None
+        return copy.deepcopy(value)
+
+    def resolve_hidden(node_id, name, default):
+        value = source_nodes[node_id].get("inputs", {}).get(name, default)
+        if _is_link(value):
+            raise ValueError(f"{name} on distributed node {node_id!r} must be a literal")
+        return copy.deepcopy(value)
+
+    resolved_distributed = {}
+
+    def resolve_distributed(node_id, output_slot):
+        node_id = str(node_id)
+        class_type = source_nodes[node_id]["class_type"]
+        if output_slot != 0:
+            raise ValueError(f"invalid output slot {output_slot} for {class_type} {node_id!r}")
+        cache_key = (node_id, output_slot)
+        if cache_key in resolved_distributed:
+            return copy.deepcopy(resolved_distributed[cache_key])
+        inputs = source_nodes[node_id].get("inputs", {})
+        if class_type == "DistributedSeed":
+            seed = resolve_literal_or_primitive(inputs.get("seed", _SEED_DEFAULT), f"seed on {node_id!r}")
+            is_worker = resolve_hidden(node_id, "is_worker", False)
+            worker_id = resolve_hidden(node_id, "worker_id", "")
+            if is_worker:
+                worker_index = _worker_index(worker_id)
+                if worker_index is not None:
+                    seed = seed + worker_index + 1
+            result = seed
+        elif class_type == "DistributedValue":
+            default_value = resolve_literal_or_primitive(
+                inputs.get("default_value", ""), f"default_value on {node_id!r}"
+            )
+            worker_values = resolve_literal_or_primitive(
+                inputs.get("worker_values", "{}"), f"worker_values on {node_id!r}"
+            )
+            is_worker = resolve_hidden(node_id, "is_worker", False)
+            worker_id = resolve_hidden(node_id, "worker_id", "")
+            try:
+                values = json.loads(worker_values) if isinstance(worker_values, str) else worker_values
+            except json.JSONDecodeError:
+                values = {}
+            if not isinstance(values, dict):
+                values = {}
+            value_type = values.get("_type", "STRING")
+            fallback = _coerce_safe(default_value, value_type)
+            result = fallback
+            if is_worker:
+                worker_index = _worker_index(worker_id)
+                if worker_index is not None:
+                    raw_value = values.get(str(worker_index + 1), "")
+                    if raw_value:
+                        try:
+                            result = _coerce(raw_value, value_type)
+                        except (ValueError, IndexError):
+                            result = fallback
+        elif class_type == "DistributedModelName":
+            text = resolve_literal_or_primitive(
+                inputs.get("text", ""), f"DistributedModelName text on {node_id!r}"
+            )
+            result = _model_name_result(text)
+        else:
+            raise ValueError(f"unsupported distributed node: {class_type}")
+        resolved_distributed[cache_key] = copy.deepcopy(result)
+        return result
+
+    return resolve_distributed
+
+
+def _build_savers(source_nodes, collector_ids):
+    """Construct stable SaveImage nodes and their collector metadata."""
+    # The prepared collector values define the artifact identity. The fallback
+    # for absent hidden values is stable and retains the collector's job scope.
+    collector_specs = []
+    max_numeric_id = 0
+    for node_id in source_nodes:
+        try:
+            max_numeric_id = max(max_numeric_id, int(node_id))
+        except ValueError:
+            pass
+    next_numeric_id = max_numeric_id + 1
+    for collector_id in collector_ids:
+        inputs = source_nodes[collector_id].get("inputs", {})
+        raw_job_id = inputs.get("multi_job_id", collector_id)
+        if _is_link(raw_job_id):
+            raise ValueError(f"collector {collector_id!r} multi_job_id must be a literal")
+        job_id = str(collector_id if raw_job_id in (None, "") else raw_job_id)
+        raw_worker_id = inputs.get("worker_id", "")
+        if _is_link(raw_worker_id):
+            raise ValueError(f"collector {collector_id!r} worker_id must be a literal")
+        worker_id = "" if raw_worker_id is None else str(raw_worker_id)
+        saver_id = str(next_numeric_id)
+        next_numeric_id += 1
+        prefix = _filename_prefix(job_id, worker_id, collector_id, "image")
+        collector_specs.append(
+            {
+                "collector_id": collector_id,
+                "job_id": job_id,
+                "node_id": saver_id,
+                "kind": "image",
+                "prefix": prefix,
+                "source_link": copy.deepcopy(inputs["images"]),
+            }
+        )
+
+    saver_nodes = {}
+    for spec in collector_specs:
+        saver_nodes[spec["node_id"]] = {
+            "class_type": "SaveImage",
+            "inputs": {
+                "images": spec["source_link"],
+                "filename_prefix": spec["prefix"],
+            },
+            "_meta": {"title": f"Serverless output: collector {spec['collector_id']} (image)"},
+        }
+
+    return collector_specs, saver_nodes
+
+
+def _prune_graph(source_nodes, collector_specs, saver_nodes, resolve_distributed):
+    """Keep saver ancestors, replacing distributed scalar links with literals."""
+    all_nodes = {node_id: copy.deepcopy(node) for node_id, node in source_nodes.items()}
+    all_nodes.update(copy.deepcopy(saver_nodes))
+    needed = set()
+    pending = [spec["node_id"] for spec in collector_specs]
+    while pending:
+        node_id = pending.pop()
+        if node_id in needed:
+            continue
+        needed.add(node_id)
+        node = all_nodes[node_id]
+        if node["class_type"] == "DistributedCollector":
+            raise ValueError("collector-to-collector dependence is unsupported")
+        for input_name, value in list(node.get("inputs", {}).items()):
+            if not _is_link(value):
+                continue
+            upstream_id, output_slot = str(value[0]), value[1]
+            upstream_class = source_nodes[upstream_id]["class_type"]
+            if upstream_class in _REMOVED_SCALAR_NODES:
+                node["inputs"][input_name] = resolve_distributed(upstream_id, output_slot)
+            else:
+                pending.append(upstream_id)
+
+    compiled = {
+        node_id: copy.deepcopy(node)
+        for node_id, node in all_nodes.items()
+        if node_id in needed
+    }
+    if any("distributed" in node["class_type"].lower() for node in compiled.values()):
+        raise ValueError("compiled prompt still contains a distributed dependency")
+
+    return compiled
+
+
+def compile_serverless_prompt(prepared_prompt: dict) -> tuple[dict, list[dict]]:
+    """Compile a prepared worker prompt to an image-only serverless graph.
+
+    Every DistributedCollector's original ``images`` input is connected to a
+    new SaveImage node. Only those savers and their upstream dependencies are
+    retained. DistributedSeed/Value/ModelName links are resolved as scalar
+    literals; unsupported node execution is never attempted.
+    """
+    source_nodes, collector_ids = _validate_prepared_prompt(prepared_prompt)
+    resolve_distributed = _make_scalar_resolver(source_nodes)
+    collector_specs, saver_nodes = _build_savers(source_nodes, collector_ids)
+    compiled = _prune_graph(
+        source_nodes,
+        collector_specs,
+        saver_nodes,
+        resolve_distributed,
+    )
+    manifest = [
+        {key: spec[key] for key in ("collector_id", "job_id", "node_id", "kind", "prefix")}
+        for spec in collector_specs
+    ]
+    return compiled, manifest

api/orchestration/serverless_runtime.py

@@ -0,0 +1,380 @@
+"""Master-owned asynchronous Runpod jobs, independent of worker callbacks."""
+import asyncio
+import time
+import uuid
+from pathlib import Path
+
+from .runpod_client import RunpodClient, RunpodError, AmbiguousSubmission
+from .serverless_ledger import JobLedger
+from .serverless_state import (
+    PROVIDER_STATES, PROVIDER_TERMINAL, belongs_to_pool, is_unresolved,
+    occupies_capacity, terminal_provider_status,
+)
+
+
+# Fixed per-worker cap on this master's outstanding Runpod jobs; equals the
+# Images per run maximum (MAX_JOBS_PER_RUN in utils/serverless_config.py), so
+# every image of a run is submitted at once. Deliberately not configurable or
+# derived from jobs_per_run: unresolved jobs from earlier runs (ambiguous
+# submissions, unconfirmed cancels) still count against it, and a constant
+# never has to change while such jobs exist.
+PER_WORKER_JOB_LIMIT = 16
+
+
+async def decode_results(output, manifest):
+    from .serverless_media import collect_images
+    import torch
+    images = await collect_images(output, manifest)
+    return {job_id: [torch.from_numpy(image).unsqueeze(0) for image in batch] for job_id, batch in images.items()}
+
+
+def state_directory():
+    import folder_paths
+    return Path(folder_paths.get_user_directory()) / "distributed" / "serverless"
+
+
+def get_secret_store():
+    from ...utils.serverless_config import SecretStore
+    return SecretStore(state_directory() / "credentials.json")
+
+
+def get_manager():
+    import server
+    from ...utils.network import get_client_session
+    ps = server.PromptServer.instance
+    manager = getattr(ps, "distributed_serverless", None)
+    if manager is None:
+        manager = ServerlessManager(ps, get_secret_store(), state_directory() / "jobs.json", get_client_session)
+        ps.distributed_serverless = manager
+    return manager
+
+
+class ServerlessManager:
+    def __init__(self, prompt_server, secrets, ledger_path, session_factory):
+        self.server = prompt_server
+        self.secrets = secrets
+        self.ledger = JobLedger(ledger_path)
+        self.session_factory = session_factory
+        self.tasks = {}
+        self.runs = {}
+        self.pool_limits = {}
+        self.health_cache = {}
+        self.health_locks = {}
+        self.poll_seconds = 2.0
+        self.capacity_changed = asyncio.Condition()
+        self.recovery_lock = asyncio.Lock()
+        self._recover_local_records()
+
+    def _recover_local_records(self):
+        # WAITING is committed before admission; it cannot have reached Runpod.
+        # SUBMITTING is an intent record, so a lost ID must remain reserved.
+        for record in self.ledger.all():
+            if record.get("status") == "WAITING" and not record.get("job_id"):
+                self.ledger.update(record["record_id"], status="CANCELLED", updated_at=time.time())
+            elif record.get("status") == "SUBMITTING" and not record.get("job_id"):
+                self.ledger.update(record["record_id"], status="UNKNOWN", updated_at=time.time())
+
+    def _task_is_active(self, record_id):
+        task = self.tasks.get(record_id)
+        return task is not None and not task.done()
+
+    def pool_is_active(self, worker_id):
+        """Protect both live local work and unresolved pre-restart provider work."""
+        for rid, run in self.runs.items():
+            pool_id = run["worker"].get("pool_id", run["worker"]["id"])
+            if str(pool_id) == str(worker_id) and self._task_is_active(rid):
+                return True
+        return any(belongs_to_pool(record, worker_id) and is_unresolved(record)
+                   for record in self.ledger.all())
+
+    async def _reserve_provider_capacity(self, run):
+        pool_id = run["worker"].get("pool_id", run["worker"]["id"])
+        async with self.capacity_changed:
+            while True:
+                self._check_master(run)
+                outstanding = sum(belongs_to_pool(record, pool_id) and occupies_capacity(record)
+                                  for record in self.ledger.all())
+                if outstanding < PER_WORKER_JOB_LIMIT:
+                    # Claim atomically before any await that could submit a billed job.
+                    self.ledger.update(run["record_id"], status="SUBMITTING", updated_at=time.time())
+                    return
+                await self.capacity_changed.wait()
+
+    async def _notify_capacity(self):
+        async with self.capacity_changed:
+            self.capacity_changed.notify_all()
+
+    async def _client(self, worker):
+        pool_id = worker.get("pool_id", worker["id"])
+        key = self.secrets.get(pool_id)
+        return RunpodClient(await self.session_factory(), worker["endpoint_id"], key)
+
+    async def health(self, worker, *, fresh=False):
+        pool_id = worker.get("pool_id", worker["id"])
+        cache_key = (pool_id, worker["endpoint_id"])
+        lock = self.health_locks.setdefault(cache_key, asyncio.Lock())
+        async with lock:
+            cached = self.health_cache.get(cache_key)
+            if cached and not fresh and time.monotonic() - cached[0] < 10:
+                return cached[1]
+            client = await self._client(worker)
+            data = await client.health()
+            jobs = data.get("jobs", {})
+            if not isinstance(jobs, dict) or not isinstance(data.get("workers"), dict):
+                raise RunpodError("Runpod returned an invalid endpoint health record")
+            try:
+                pending = max(0, int(jobs.get("inQueue", 0))) + max(0, int(jobs.get("inProgress", 0)))
+            except (ValueError, TypeError):
+                raise RunpodError("Runpod returned invalid queue counts") from None
+            result = {"ok": True, "queueRemaining": pending, "message": "Queued / processing" if pending else "Ready (scales from zero)"}
+            self.health_cache[cache_key] = (time.monotonic(), result)
+            return result
+
+    def _semaphore(self, worker):
+        # The limit is a constant, so it never has to be swapped while this
+        # worker still has unresolved jobs.
+        pool_id = worker.get("pool_id", worker["id"])
+        semaphore = self.pool_limits.get(pool_id)
+        if semaphore is None:
+            semaphore = asyncio.Semaphore(PER_WORKER_JOB_LIMIT)
+            self.pool_limits[pool_id] = semaphore
+        return semaphore
+
+    def ensure_capacity(self, count):
+        if sum(not task.done() for task in self.tasks.values()) + count > 64:
+            raise ValueError("Distributed already has too many pending serverless jobs (limit 64)")
+
+    async def start(self, execution_id, prompt_id, prepared_jobs):
+        self.ensure_capacity(len(prepared_jobs))
+        new_runs = []
+        # Journal every local slot before starting any billed submission.
+        for prepared in prepared_jobs:
+            worker = prepared["worker"]
+            semaphore = self._semaphore(worker)
+            record_id = uuid.uuid4().hex
+            now = time.time()
+            self.ledger.update(record_id, execution_id=execution_id, prompt_id=prompt_id,
+                               worker_id=worker.get("pool_id", worker["id"]), participant_id=worker["id"],
+                               endpoint_id=worker["endpoint_id"], job_id=None, status="WAITING", created_at=now, updated_at=now)
+            run = dict(prepared, execution_id=execution_id, prompt_id=prompt_id, record_id=record_id,
+                       provider_job_id=None, client=None, deadline=time.monotonic() + worker["timeout_seconds"])
+            self.runs[record_id] = run
+            new_runs.append((record_id, run, semaphore))
+        for record_id, run, semaphore in new_runs:
+            self.tasks[record_id] = asyncio.create_task(self._run(run, semaphore), name=f"distributed-serverless-{record_id}")
+        # Bound in-memory completed tracking as well as the on-disk journal.
+        done = [rid for rid, task in self.tasks.items() if task.done()]
+        for rid in done[:-200]:
+            self.tasks.pop(rid, None)
+            self.runs.pop(rid, None)
+
+    def is_pending(self, collector_job_id, participant_id):
+        return any(not task.done() and run["worker"]["id"] == participant_id and
+                   any(m["job_id"] == collector_job_id for m in run["manifest"])
+                   for rid, task in self.tasks.items() for run in [self.runs[rid]])
+
+    def owns(self, collector_job_id, participant_id):
+        return any(run["worker"]["id"] == participant_id and
+                   any(m["job_id"] == collector_job_id for m in run["manifest"]) for run in self.runs.values())
+
+    async def _deliver(self, run, images=None, error=None):
+        async with self.server.distributed_jobs_lock:
+            for binding in run["manifest"]:
+                queue = self.server.distributed_pending_jobs.get(binding["job_id"])
+                if queue is None:
+                    continue  # A cancelled/finished collector must never be resurrected.
+                if error:
+                    queue.put_nowait({"worker_id": run["worker"]["id"], "serverless_error": error, "is_last": True})
+                else:
+                    batch = images[binding["job_id"]]
+                    for index, tensor in enumerate(batch):
+                        queue.put_nowait({"worker_id": run["worker"]["id"], "image_index": index,
+                                          "tensor": tensor, "is_last": index == len(batch) - 1})
+
+    async def _cancel_provider(self, run):
+        if run["provider_job_id"] and run["client"]:
+            try:
+                result = await run["client"].cancel(run["provider_job_id"])
+                return terminal_provider_status(result)
+            except Exception:
+                return None
+        return None
+
+    def _check_master(self, run):
+        queue = getattr(self.server, "prompt_queue", None)
+        if queue is None:
+            return
+        history = queue.get_history(run["prompt_id"])
+        record = history.get(run["prompt_id"], {})
+        if record.get("status", {}).get("status_str") == "error":
+            raise RunpodError("Master workflow failed; serverless work stopped")
+        if not record and hasattr(queue, "get_current_queue"):
+            running, pending = queue.get_current_queue()
+            if not any(str(item[1]) == run["prompt_id"] for item in [*running, *pending]):
+                raise RunpodError("Master workflow was removed from the queue; serverless work stopped")
+
+    async def _run(self, run, semaphore):
+        rid = run["record_id"]
+        submitted = False
+        provider_terminal = False
+        try:
+            async with asyncio.timeout(max(0, run["deadline"] - time.monotonic())):
+                async with semaphore:
+                    await self._reserve_provider_capacity(run)
+                    client = await self._client(run["worker"])
+                    run["client"] = client
+                    submitted = True
+                    job_id = await client.submit(run["workflow"], run["images"], run["worker"]["timeout_seconds"])
+                    run["provider_job_id"] = job_id
+                    # Persist immediately: all subsequent actions use this exact ID.
+                    self.ledger.update(rid, job_id=job_id, status="IN_QUEUE", updated_at=time.time())
+                    while True:
+                        self._check_master(run)
+                        response = await client.status(job_id)
+                        status = response.get("status")
+                        if not isinstance(status, str) or status not in PROVIDER_STATES:
+                            raise RunpodError("Runpod returned an unknown job state")
+                        provider_terminal = status in PROVIDER_TERMINAL
+                        if status == "COMPLETED":
+                            output = response.get("output")
+                            if not isinstance(output, dict) or output.get("error") or output.get("errors"):
+                                raise RunpodError("Serverless workflow reported an output error; inspect the provider job")
+                            self.ledger.update(rid, status="DOWNLOADING", updated_at=time.time())
+                            images = await decode_results(output, run["manifest"])
+                            await self._deliver(run, images=images)
+                            self.ledger.update(rid, status="COMPLETED", updated_at=time.time())
+                            return
+                        self.ledger.update(rid, status=status, updated_at=time.time())
+                        if provider_terminal:
+                            raise RunpodError(f"Serverless job ended with {status}")
+                        await asyncio.sleep(self.poll_seconds)
+        except AmbiguousSubmission as exc:
+            self.ledger.update(rid, status="UNKNOWN", error=str(exc), updated_at=time.time())
+            await self._deliver(run, error=str(exc))
+        except (asyncio.TimeoutError, asyncio.CancelledError) as exc:
+            terminal = "FAILED" if provider_terminal else await self._cancel_provider(run)
+            if submitted and not run["provider_job_id"]:
+                status = "UNKNOWN"
+                error = "Submission interrupted before its job ID was received; check Runpod before retrying"
+            elif isinstance(exc, asyncio.TimeoutError):
+                status = terminal or ("CANCEL_REQUESTED" if run["provider_job_id"] else "TIMED_OUT")
+                error = "Serverless job exceeded its queue/execution deadline"
+            else:
+                status = terminal or ("CANCEL_REQUESTED" if run["provider_job_id"] else "CANCELLED")
+                error = "Serverless work stopped" if terminal or not submitted else "Serverless cancellation requested; provider confirmation is pending"
+            self.ledger.update(rid, status=status, error=error, updated_at=time.time())
+            await self._deliver(run, error=error)
+        except Exception as exc:
+            # If polling/decoding failed, stop only our own provider job, never the pool.
+            terminal = "FAILED" if provider_terminal else await self._cancel_provider(run)
+            error = str(exc) if isinstance(exc, RunpodError) else "Serverless processing failed; check the workflow and endpoint contract"
+            status = terminal or ("CANCEL_REQUESTED" if run["provider_job_id"] else "FAILED")
+            self.ledger.update(rid, status=status, error=error, updated_at=time.time())
+            await self._deliver(run, error=error)
+        finally:
+            # Completed tracking retains IDs/bindings only, not potentially large
+            # base64 inputs, workflow metadata, or a client holding credentials.
+            run.pop("workflow", None)
+            run.pop("images", None)
+            run["client"] = None
+            # A released local semaphore alone cannot release a durable reservation.
+            await self._notify_capacity()
+
+    async def wait_for_execution(self, execution_id):
+        tasks = [task for rid, task in self.tasks.items() if self.runs[rid]["execution_id"] == execution_id]
+        await asyncio.gather(*tasks, return_exceptions=False)
+
+    def _detached_records(self, worker_id=None, execution_id=None):
+        return [record for record in self.ledger.all()
+                if is_unresolved(record) and not self._task_is_active(record["record_id"])
+                and (worker_id is None or belongs_to_pool(record, worker_id))
+                and (execution_id is None or record.get("execution_id") == execution_id)]
+
+    async def _recover_provider_record(self, record, *, cancel=False):
+        """Reconcile an exact saved ID, never submit or resume a workflow."""
+        if not record.get("job_id"):
+            return False
+        status = None
+        try:
+            client = await self._client({"id": record["worker_id"], "endpoint_id": record["endpoint_id"]})
+            response = await (client.cancel(record["job_id"]) if cancel else client.status(record["job_id"]))
+            status = terminal_provider_status(response)
+        except Exception:
+            pass  # Unreachable/expired is not proof of provider termination.
+        if status is not None or cancel:
+            self.ledger.update(record["record_id"], status=status or "CANCEL_REQUESTED", updated_at=time.time())
+        return status is not None
+
+    async def reconcile(self, worker_id=None):
+        """Read terminal states for detached jobs and release only proven capacity."""
+        resolved = 0
+        async with self.recovery_lock:
+            for record in self._detached_records(worker_id):
+                resolved += await self._recover_provider_record(record)
+        await self._notify_capacity()
+        return resolved
+
+    async def acknowledge_recovery(self, worker_id, record_id, *, provider_checked=False):
+        """Explicit operator escape hatch for lost IDs/expired provider records.
+
+        This records a human's verification, not a cancellation or provider
+        status. It is never called automatically or allowed for live tasks.
+        """
+        if provider_checked is not True:
+            raise ValueError("Check Runpod and confirm that this job is no longer running before resolving it")
+        async with self.recovery_lock:
+            record = next((r for r in self._detached_records(worker_id) if r["record_id"] == record_id), None)
+            if record is None:
+                raise ValueError("Unresolved inactive job not found for this worker")
+            self.ledger.update(record_id, status="RESOLVED", updated_at=time.time(),
+                               error="Recovery explicitly acknowledged after operator checked Runpod; not provider-confirmed")
+        await self._notify_capacity()
+
+    async def cancel(self, worker_id=None, execution_id=None, collector_job_id=None):
+        targets = {}
+        for rid, task in self.tasks.items():
+            run = self.runs[rid]
+            if worker_id and run["worker"].get("pool_id", run["worker"]["id"]) != worker_id:
+                continue
+            if execution_id and run["execution_id"] != execution_id:
+                continue
+            if collector_job_id and not any(m["job_id"] == collector_job_id for m in run["manifest"]):
+                continue
+            if not task.done():
+                # Repeated requests must not interrupt the first cancellation's cleanup.
+                if not task.cancelling():
+                    task.cancel()
+                targets[rid] = task
+        if targets:
+            await asyncio.gather(*targets.values(), return_exceptions=True)
+        recovered = 0
+        async with self.recovery_lock:
+            for record in self._detached_records(worker_id, execution_id):
+                rid = record["record_id"]
+                # A task cancelled before its first step never enters _run's handlers.
+                # A failed start() may also leave a journalled but unstarted local slot.
+                if record.get("status") == "WAITING" and not record.get("job_id"):
+                    run = self.runs.get(rid)
+                    if run is not None and (collector_job_id is None or any(m["job_id"] == collector_job_id for m in run["manifest"])):
+                        self.ledger.update(rid, status="CANCELLED", updated_at=time.time())
+                        run.pop("workflow", None)
+                        run.pop("images", None)
+                        await self._deliver(run, error="Serverless job cancelled before submission")
+                    continue
+                # Do not send two cancellation requests for the same live job in
+                # one action. Its unresolved reservation survives for later recovery.
+                if rid in targets or collector_job_id or not record.get("job_id"):
+                    continue
+                await self._recover_provider_record(record, cancel=True)
+                recovered += 1
+        await self._notify_capacity()
+        return len(targets) + recovered
+
+    def public_jobs(self, worker_id):
+        records = [record for record in reversed(self.ledger.all()) if belongs_to_pool(record, worker_id)]
+        # Never hide unresolved IDs behind the bounded completed history.
+        visible = [record for record in records if is_unresolved(record)]
+        visible.extend(record for record in records if not is_unresolved(record))
+        unresolved_count = sum(is_unresolved(record) for record in records)
+        return [{key: record.get(key) for key in ("record_id", "job_id", "status", "worker_id", "execution_id", "error")}
+                for record in visible[:max(20, unresolved_count)]]

api/orchestration/serverless_state.py

@@ -0,0 +1,29 @@
+"""Shared rules for provider capacity and durable recovery protection.
+
+A local task ending is not evidence that provider work ended. RESOLVED is an
+explicit operator acknowledgment, never a provider-reported state.
+"""
+
+PROVIDER_TERMINAL = frozenset({"COMPLETED", "FAILED", "CANCELLED", "TIMED_OUT"})
+TERMINAL = PROVIDER_TERMINAL | {"RESOLVED"}
+PROVIDER_STATES = PROVIDER_TERMINAL | {"IN_QUEUE", "IN_PROGRESS"}
+
+
+def terminal_provider_status(response):
+    status = response.get("status") if isinstance(response, dict) else None
+    return status if isinstance(status, str) and status in PROVIDER_TERMINAL else None
+
+
+def is_unresolved(record):
+    return record.get("status") not in TERMINAL
+
+
+def occupies_capacity(record):
+    # WAITING is journalled before admission and cannot have been submitted.
+    # Every other unresolved state is conservatively reserved, including a
+    # submission whose provider ID was lost and an unconfirmed cancellation.
+    return is_unresolved(record) and (record.get("status") != "WAITING" or bool(record.get("job_id")))
+
+
+def belongs_to_pool(record, worker_id):
+    return str(record.get("worker_id")) == str(worker_id)

api/queue_orchestration.py

@@ -1,7 +1,8 @@
 import asyncio
 import time
 import uuid
-from typing import Any
+
+import server
 
 from ..utils.async_helpers import queue_prompt_payload
 from ..utils.config import load_config
@@ -12,17 +13,20 @@ from ..utils.constants import (
     ORCHESTRATION_WORKER_PREP_CONCURRENCY,
 )
 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.network import build_master_url, build_master_callback_url
 from ..utils.trace_logger import trace_debug
-from .schemas import coerce_positive_float, coerce_positive_int
+from .schemas import parse_positive_float, parse_positive_int
 from .orchestration.dispatch import (
     dispatch_worker_prompt,
-    rank_workers_by_load,
     select_active_workers,
     select_least_busy_worker,
 )
-from .orchestration.media_sync import convert_paths_for_platform, fetch_worker_path_separator, sync_worker_media
+from .orchestration.media_sync import (
+    convert_paths_for_platform,
+    fetch_worker_has_node,
+    fetch_worker_path_separator,
+    sync_worker_media,
+)
 from .orchestration.prompt_transform import (
     PromptIndex,
     apply_participant_overrides,
@@ -30,23 +34,36 @@ from .orchestration.prompt_transform import (
     generate_job_id_map,
     prepare_delegate_master_prompt,
     prune_prompt_for_worker,
+    WORKER_SINK_CLASS,
 )
 
+
+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."""
-    ensure_distributed_runtime_state(server_instance)
+    ps = server_instance or prompt_server
+    if not hasattr(ps, "distributed_pending_jobs"):
+        ps.distributed_pending_jobs = {}
+    if not hasattr(ps, "distributed_jobs_lock"):
+        ps.distributed_jobs_lock = asyncio.Lock()
+
+
+# Initialize top-level distributed queue state at module import time.
+ensure_distributed_state()
 
 
 async def _ensure_distributed_queue(job_id):
     """Ensure a queue exists for the given distributed job ID."""
-    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()
+    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()
 
 
 def _resolve_enabled_workers(config, requested_ids=None):
@@ -63,6 +80,14 @@ def _resolve_enabled_workers(config, requested_ids=None):
         elif not worker.get("enabled", False):
             continue
 
+        if worker.get("type") == "serverless":
+            from ..utils.serverless_config import validate_serverless_worker
+            pool = validate_serverless_worker(worker)
+            for slot in range(pool["jobs_per_run"]):
+                workers.append(dict(pool, pool_id=worker_id, id=f"serverless:{worker_id}:{slot + 1}",
+                                    name=f"{pool.get('name', worker_id)} / job {slot + 1}"))
+            continue
+
         raw_port = worker.get("port", worker.get("listen_port", 8188))
         try:
             port = int(raw_port or 8188)
@@ -85,19 +110,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 = coerce_positive_int(
+    worker_probe_concurrency = parse_positive_int(
         settings.get("worker_probe_concurrency"),
         ORCHESTRATION_WORKER_PROBE_CONCURRENCY,
     )
-    worker_prep_concurrency = coerce_positive_int(
+    worker_prep_concurrency = parse_positive_int(
         settings.get("worker_prep_concurrency"),
         ORCHESTRATION_WORKER_PREP_CONCURRENCY,
     )
-    media_sync_concurrency = coerce_positive_int(
+    media_sync_concurrency = parse_positive_int(
         settings.get("media_sync_concurrency"),
         ORCHESTRATION_MEDIA_SYNC_CONCURRENCY,
     )
-    media_sync_timeout_seconds = coerce_positive_float(
+    media_sync_timeout_seconds = parse_positive_float(
         settings.get("media_sync_timeout_seconds"),
         ORCHESTRATION_MEDIA_SYNC_TIMEOUT,
     )
@@ -133,6 +158,7 @@ async def _prepare_worker_payload(
     enabled_ids,
     job_id_map,
     master_url,
+    config,
     delegate_master,
     trace_execution_id,
     worker_prep_semaphore,
@@ -142,21 +168,50 @@ async def _prepare_worker_payload(
     """Prepare one worker prompt payload with bounded concurrency and media-sync timeout."""
     async with worker_prep_semaphore:
         worker_prompt = prompt_index.copy_prompt()
+        if worker.get("type") == "serverless":
+            from .orchestration.serverless_prompt import compile_serverless_prompt
+            from .orchestration.serverless_media import prepare_images
+            from .orchestration.runpod_client import serialize_request
+            worker_prompt = apply_participant_overrides(
+                prune_prompt_for_worker(worker_prompt), worker["id"], enabled_ids,
+                job_id_map, "", delegate_master, prompt_index,
+            )
+            workflow, manifest = compile_serverless_prompt(worker_prompt)
+            images = await prepare_images(workflow)
+            # Fail during preparation, before regular workers or the master are
+            # dispatched. submit() uses this same serializer for the sent bytes.
+            serialize_request(workflow, images, worker["timeout_seconds"])
+            return worker, {"worker": worker, "workflow": workflow, "manifest": manifest, "images": images}
+
+        worker_master_url = build_master_callback_url(
+            worker,
+            config=config,
+            prompt_server_instance=prompt_server,
+        )
 
         worker_type = str(worker.get("type") or "local").strip().lower()
         is_remote_like = bool(worker.get("host")) and worker_type != "local"
-        if is_remote_like:
-            path_separator = await fetch_worker_path_separator(worker, trace_execution_id=trace_execution_id)
-            if path_separator:
-                worker_prompt = convert_paths_for_platform(worker_prompt, path_separator)
+        # Older workers lack the sink node; they keep the PreviewImage terminator.
+        has_sink_task = asyncio.ensure_future(
+            fetch_worker_has_node(worker, WORKER_SINK_CLASS, trace_execution_id=trace_execution_id)
+        )
+        try:
+            if is_remote_like:
+                path_separator = await fetch_worker_path_separator(worker, trace_execution_id=trace_execution_id)
+                if path_separator:
+                    worker_prompt = convert_paths_for_platform(worker_prompt, path_separator)
+            use_worker_sink = await has_sink_task
+        finally:
+            if not has_sink_task.done():
+                has_sink_task.cancel()
 
-        worker_prompt = prune_prompt_for_worker(worker_prompt)
+        worker_prompt = prune_prompt_for_worker(worker_prompt, use_worker_sink=use_worker_sink)
         worker_prompt = apply_participant_overrides(
             worker_prompt,
             worker["id"],
             enabled_ids,
             job_id_map,
-            master_url,
+            worker_master_url,
             delegate_master,
             prompt_index,
         )
@@ -180,16 +235,60 @@ async def _prepare_worker_payload(
         return worker, worker_prompt
 
 
-async def _select_execution_workers(
-    workers,
-    use_websocket,
-    delegate_master,
-    load_balance_requested,
-    has_branch_nodes,
-    master_url,
-    execution_trace_id,
-    worker_probe_concurrency,
+async def orchestrate_distributed_execution(
+    prompt_obj,
+    workflow_meta,
+    client_id,
+    enabled_worker_ids=None,
+    delegate_master=None,
+    trace_execution_id=None,
 ):
+    """Core orchestration logic for the /distributed/queue endpoint.
+
+    Returns:
+        tuple[str, int, int, dict]: (prompt_id, number, worker_count, node_errors)
+    """
+    ensure_distributed_state()
+    execution_trace_id = trace_execution_id or _generate_execution_trace_id()
+
+    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)
+    (
+        worker_probe_concurrency,
+        worker_prep_concurrency,
+        media_sync_concurrency,
+        media_sync_timeout_seconds,
+    ) = _resolve_orchestration_limits(config)
+    requested_ids = enabled_worker_ids if enabled_worker_ids is not None else None
+    workers = _resolve_enabled_workers(config, requested_ids)
+    prompt_index = PromptIndex(prompt_obj)
+    load_balance_requested = _prompt_requests_load_balance(prompt_index)
+    trace_debug(
+        execution_trace_id,
+        (
+            f"Orchestration start: requested_workers={len(workers)}, "
+            f"requested_ids={requested_ids if requested_ids is not None else 'enabled_only'}, "
+            f"websocket={use_websocket}, "
+            f"probe_concurrency={worker_probe_concurrency}, "
+            f"prep_concurrency={worker_prep_concurrency}, "
+            f"media_sync_concurrency={media_sync_concurrency}, "
+            f"media_sync_timeout={media_sync_timeout_seconds:.1f}s, "
+            f"load_balance={load_balance_requested}"
+        ),
+    )
+
+    # Respect master delegate-only configuration
+    if delegate_master is None:
+        delegate_master = bool(config.get("settings", {}).get("master_delegate_only", False))
+
+    if not workers and delegate_master:
+        trace_debug(
+            execution_trace_id,
+            "Delegate-only requested but no workers are enabled. Falling back to master execution.",
+        )
+        delegate_master = False
+
     active_workers, delegate_master = await select_active_workers(
         workers,
         use_websocket,
@@ -201,6 +300,7 @@ async def _select_execution_workers(
     if load_balance_requested:
         candidate_workers = list(active_workers)
         if not delegate_master:
+            # Include master in load balancing only when master participation is enabled.
             candidate_workers.append(
                 {
                     "id": "master",
@@ -225,6 +325,7 @@ async def _select_execution_workers(
             selected_worker = candidate_workers[0]
 
         if selected_worker is not None and str(selected_worker.get("id")) == "master":
+            # Master selected as least busy; run master workload only.
             active_workers = []
             delegate_master = False
             trace_debug(
@@ -233,6 +334,7 @@ async def _select_execution_workers(
             )
         elif selected_worker is not None:
             active_workers = [selected_worker]
+            # Worker selected as least busy; keep master orchestrator-only for this run.
             delegate_master = True
             trace_debug(
                 execution_trace_id,
@@ -246,47 +348,26 @@ async def _select_execution_workers(
             active_workers = []
             delegate_master = False
 
-    if has_branch_nodes and len(active_workers) > 1:
-        active_workers = await rank_workers_by_load(
-            active_workers,
-            trace_execution_id=execution_trace_id,
-            probe_concurrency=worker_probe_concurrency,
+    enabled_ids = [worker["id"] for worker in active_workers]
+
+    discovery_prefix = f"exec_{int(time.time() * 1000)}_{uuid.uuid4().hex[:6]}"
+    job_id_map = generate_job_id_map(prompt_index, discovery_prefix)
+
+    if not job_id_map:
+        trace_debug(execution_trace_id, "No distributed nodes detected; queueing prompt on master only.")
+        queue_result = await queue_prompt_payload(
+            prompt_obj,
+            workflow_meta,
+            client_id,
+            include_queue_metadata=True,
+        )
+        return (
+            queue_result["prompt_id"],
+            queue_result["number"],
+            0,
+            queue_result.get("node_errors", {}),
         )
 
-    return active_workers, delegate_master
-
-
-def _distributed_queue_nodes(prompt_index):
-    return (
-        prompt_index.nodes_for_class("DistributedCollector")
-        + 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)
-        if job_id:
-            await _ensure_distributed_queue(job_id)
-
-
-def _build_master_prompt(
-    prompt_index,
-    enabled_ids,
-    job_id_map,
-    master_url,
-    delegate_master,
-):
     master_prompt = prompt_index.copy_prompt()
     master_prompt = apply_participant_overrides(
         master_prompt,
@@ -298,171 +379,19 @@ def _build_master_prompt(
         prompt_index,
     )
 
-    if not delegate_master:
-        return master_prompt
-
-    collector_ids = (
-        find_nodes_by_class(master_prompt, "DistributedCollector")
-        + find_nodes_by_class(master_prompt, "DistributedBranchCollector")
-    )
-    upscale_nodes = find_nodes_by_class(master_prompt, "UltimateSDUpscaleDistributed")
-    # Include USDU nodes as collector-like for delegate pruning
-    collector_ids.extend(upscale_nodes)
-    if not collector_ids:
-        debug_log(
-            "Delegate-only master mode requested but no collector/branch-collector nodes found in master prompt. Running full prompt on master."
-        )
-        return master_prompt
-    return prepare_delegate_master_prompt(master_prompt, collector_ids)
-
-
-async def _prepare_worker_payloads_for_dispatch(
-    active_workers,
-    prompt_index,
-    enabled_ids,
-    job_id_map,
-    master_url,
-    delegate_master,
-    execution_trace_id,
-    worker_prep_concurrency,
-    media_sync_concurrency,
-    media_sync_timeout_seconds,
-):
-    if not active_workers:
-        return []
-
-    worker_prep_semaphore = asyncio.Semaphore(worker_prep_concurrency)
-    media_sync_semaphore = asyncio.Semaphore(media_sync_concurrency)
-    return await asyncio.gather(
-        *[
-            _prepare_worker_payload(
-                worker,
-                prompt_index,
-                enabled_ids,
-                job_id_map,
-                master_url,
-                delegate_master,
-                execution_trace_id,
-                worker_prep_semaphore,
-                media_sync_semaphore,
-                media_sync_timeout_seconds,
+    if delegate_master:
+        collector_ids = find_nodes_by_class(master_prompt, "DistributedCollector")
+        upscale_nodes = find_nodes_by_class(master_prompt, "UltimateSDUpscaleDistributed")
+        if upscale_nodes:
+            debug_log(
+                "Delegate-only master mode currently does not support UltimateSDUpscaleDistributed nodes; running full prompt on master."
             )
-            for worker in active_workers
-        ]
-    )
-
-
-async def _dispatch_worker_payloads(
-    worker_payloads,
-    workflow_meta,
-    client_id,
-    use_websocket,
-    execution_trace_id,
-):
-    if not worker_payloads:
-        return
-
-    await asyncio.gather(
-        *[
-            dispatch_worker_prompt(
-                worker,
-                wprompt,
-                workflow_meta,
-                client_id,
-                use_websocket=use_websocket,
-                trace_execution_id=execution_trace_id,
+        elif not collector_ids:
+            debug_log(
+                "Delegate-only master mode requested but no collectors found in master prompt. Running full prompt on master."
             )
-            for worker, wprompt in worker_payloads
-        ]
-    )
-
-
-async def orchestrate_distributed_execution(
-    prompt_obj: dict[str, Any],
-    workflow_meta: dict[str, Any] | None,
-    client_id: str | None,
-    enabled_worker_ids: list[str] | set[str] | None = None,
-    delegate_master: bool | None = None,
-    trace_execution_id: str | None = None,
-) -> tuple[str, int]:
-    """Core orchestration logic for the /distributed/queue endpoint.
-
-    Returns:
-        tuple[str, int]: (prompt_id, worker_count)
-    """
-    ensure_distributed_state()
-    execution_trace_id = trace_execution_id or _generate_execution_trace_id()
-
-    config = load_config()
-    use_websocket = bool(config.get("settings", {}).get("websocket_orchestration", False))
-    master_url = build_master_url(config=config, prompt_server_instance=get_prompt_server_instance())
-    (
-        worker_probe_concurrency,
-        worker_prep_concurrency,
-        media_sync_concurrency,
-        media_sync_timeout_seconds,
-    ) = _resolve_orchestration_limits(config)
-    requested_ids = enabled_worker_ids if enabled_worker_ids is not None else None
-    workers = _resolve_enabled_workers(config, requested_ids)
-    prompt_index = PromptIndex(prompt_obj)
-    load_balance_requested = _prompt_requests_load_balance(prompt_index)
-    has_branch_nodes = bool(prompt_index.nodes_for_class("DistributedBranch"))
-    trace_debug(
-        execution_trace_id,
-        (
-            f"Orchestration start: requested_workers={len(workers)}, "
-            f"requested_ids={requested_ids if requested_ids is not None else 'enabled_only'}, "
-            f"websocket={use_websocket}, "
-            f"probe_concurrency={worker_probe_concurrency}, "
-            f"prep_concurrency={worker_prep_concurrency}, "
-            f"media_sync_concurrency={media_sync_concurrency}, "
-            f"media_sync_timeout={media_sync_timeout_seconds:.1f}s, "
-            f"load_balance={load_balance_requested}, has_branch_nodes={has_branch_nodes}"
-        ),
-    )
-
-    # Respect master delegate-only configuration
-    if delegate_master is None:
-        delegate_master = bool(config.get("settings", {}).get("master_delegate_only", False))
-
-    if not workers and delegate_master:
-        trace_debug(
-            execution_trace_id,
-            "Delegate-only requested but no workers are enabled. Falling back to master execution.",
-        )
-        delegate_master = False
-
-    active_workers, delegate_master = await _select_execution_workers(
-        workers=workers,
-        use_websocket=use_websocket,
-        delegate_master=delegate_master,
-        load_balance_requested=load_balance_requested,
-        has_branch_nodes=has_branch_nodes,
-        master_url=master_url,
-        execution_trace_id=execution_trace_id,
-        worker_probe_concurrency=worker_probe_concurrency,
-    )
-
-    enabled_ids = [worker["id"] for worker in active_workers]
-
-    discovery_prefix = f"exec_{int(time.time() * 1000)}_{uuid.uuid4().hex[:6]}"
-    job_id_map = generate_job_id_map(prompt_index, discovery_prefix)
-
-    if not job_id_map:
-        trace_debug(execution_trace_id, "No distributed nodes detected; queueing prompt on master only.")
-        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(
-        prompt_index=prompt_index,
-        enabled_ids=enabled_ids,
-        job_id_map=job_id_map,
-        master_url=master_url,
-        delegate_master=delegate_master,
-    )
+        else:
+            master_prompt = prepare_delegate_master_prompt(master_prompt, collector_ids)
 
     if active_workers:
         trace_debug(
@@ -470,29 +399,88 @@ async def orchestrate_distributed_execution(
             "Active distributed workers: "
             + ", ".join(f"{worker['name']} ({worker['id']})" for worker in active_workers),
         )
-    worker_payloads = await _prepare_worker_payloads_for_dispatch(
-        active_workers=active_workers,
-        prompt_index=prompt_index,
-        enabled_ids=enabled_ids,
-        job_id_map=job_id_map,
-        master_url=master_url,
-        delegate_master=delegate_master,
-        execution_trace_id=execution_trace_id,
-        worker_prep_concurrency=worker_prep_concurrency,
-        media_sync_concurrency=media_sync_concurrency,
-        media_sync_timeout_seconds=media_sync_timeout_seconds,
-    )
-    await _dispatch_worker_payloads(
-        worker_payloads=worker_payloads,
-        workflow_meta=workflow_meta,
-        client_id=client_id,
-        use_websocket=use_websocket,
-        execution_trace_id=execution_trace_id,
-    )
+    worker_payloads = []
+    if active_workers:
+        worker_prep_semaphore = asyncio.Semaphore(worker_prep_concurrency)
+        media_sync_semaphore = asyncio.Semaphore(media_sync_concurrency)
+        worker_payloads = await asyncio.gather(
+            *[
+                _prepare_worker_payload(
+                    worker,
+                    prompt_index,
+                    enabled_ids,
+                    job_id_map,
+                    master_url,
+                    config,
+                    delegate_master,
+                    execution_trace_id,
+                    worker_prep_semaphore,
+                    media_sync_semaphore,
+                    media_sync_timeout_seconds,
+                )
+                for worker in active_workers
+            ]
+        )
 
-    prompt_id = await queue_prompt_payload(master_prompt, workflow_meta, client_id)
+    serverless_jobs = [payload for worker, payload in worker_payloads if worker.get("type") == "serverless"]
+    regular_payloads = [(worker, payload) for worker, payload in worker_payloads if worker.get("type") != "serverless"]
+    manager = None
+    if serverless_jobs:
+        from .orchestration.serverless_runtime import get_manager
+        manager = get_manager()
+        manager.ensure_capacity(len(serverless_jobs))
+
+    # Do not leave collector queues behind when offline compilation/media validation fails.
+    for job_id in job_id_map.values():
+        await _ensure_distributed_queue(job_id)
+
+    if regular_payloads:
+        await asyncio.gather(
+            *[
+                dispatch_worker_prompt(
+                    worker,
+                    wprompt,
+                    workflow_meta,
+                    client_id,
+                    use_websocket=use_websocket,
+                    trace_execution_id=execution_trace_id,
+                )
+                for worker, wprompt in regular_payloads
+            ]
+        )
+
+    try:
+        queue_result = await queue_prompt_payload(
+            master_prompt,
+            workflow_meta,
+            client_id,
+            include_queue_metadata=True,
+        )
+    except Exception:
+        if serverless_jobs:
+            async with prompt_server.distributed_jobs_lock:
+                for job_id in job_id_map.values():
+                    prompt_server.distributed_pending_jobs.pop(job_id, None)
+        raise
+    prompt_id = queue_result["prompt_id"]
+    if manager is not None:
+        try:
+            await manager.start(execution_trace_id, prompt_id, serverless_jobs)
+        except Exception:
+            # Wake queued collectors with an explicit error if journal/capacity setup failed.
+            await manager.cancel(execution_id=execution_trace_id)
+            async with prompt_server.distributed_jobs_lock:
+                for prepared in serverless_jobs:
+                    for binding in prepared["manifest"]:
+                        q = prompt_server.distributed_pending_jobs.get(binding["job_id"])
+                        if q is not None:
+                            q.put_nowait({"worker_id": prepared["worker"]["id"], "is_last": True,
+                                          "serverless_error": "Serverless job setup failed; no automatic retry was attempted"})
+            raise
+    prompt_number = queue_result["number"]
+    node_errors = queue_result.get("node_errors", {})
     trace_debug(
         execution_trace_id,
         f"Orchestration complete: prompt_id={prompt_id}, dispatched_workers={len(worker_payloads)}, delegate_master={delegate_master}",
     )
-    return prompt_id, len(worker_payloads)
+    return prompt_id, prompt_number, len(worker_payloads), node_errors

api/queue_request.py

@@ -1,20 +1,16 @@
 from dataclasses import dataclass
-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
+from typing import Any, Dict, List, Optional
 
 
 @dataclass(frozen=True)
 class QueueRequestPayload:
-    prompt: dict[str, Any]
-    workflow_meta: dict[str, Any] | None
+    prompt: Dict[str, Any]
+    workflow_meta: Any
     client_id: str
-    delegate_master: bool | None
-    enabled_worker_ids: list[str]
-    trace_execution_id: str | None
+    delegate_master: Optional[bool]
+    enabled_worker_ids: List[str]
+    auto_prepare: bool
+    trace_execution_id: Optional[str]
 
 
 def parse_queue_request_payload(data: Any) -> QueueRequestPayload:
@@ -22,6 +18,11 @@ 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:
@@ -34,10 +35,6 @@ 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:
@@ -54,10 +51,7 @@ 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")
-        try:
-            enabled_ids = require_enabled_worker_ids(enabled_ids_raw)
-        except ValueError as exc:
-            raise ValueError(f"enabled_worker_ids invalid: {exc}") from exc
+        enabled_ids = [str(worker_id).strip() for worker_id in enabled_ids_raw if str(worker_id).strip()]
 
     delegate_master = data.get("delegate_master")
     if delegate_master is not None and not isinstance(delegate_master, bool):
@@ -76,9 +70,10 @@ def parse_queue_request_payload(data: Any) -> QueueRequestPayload:
 
     return QueueRequestPayload(
         prompt=prompt,
-        workflow_meta=workflow_meta,
+        workflow_meta=data.get("workflow"),
         client_id=client_id,
         delegate_master=delegate_master,
         enabled_worker_ids=enabled_ids,
+        auto_prepare=auto_prepare,
         trace_execution_id=trace_execution_id,
     )

api/request_guards.py

@@ -1,15 +0,0 @@
-"""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,29 +1,3 @@
-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):
@@ -44,43 +18,26 @@ def require_fields(data: dict, *fields) -> list[str]:
     return missing
 
 
-def validate_worker_id(worker_id: str, config: dict[str, Any]) -> bool:
+def validate_worker_id(worker_id: str, config: dict) -> bool:
     """Return True when worker_id exists in config['workers']."""
-    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.")
-
+    worker_id_str = str(worker_id)
     workers = (config or {}).get("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.")
+    return any(str(worker.get("id")) == worker_id_str for worker in workers)
 
 
-def require_positive_int(value, field_name: str = "value") -> int:
-    """Strict positive-int parser for request boundary validation."""
+def validate_positive_int(value, field_name: str) -> str | None:
+    """Validate positive integers and return an error string when invalid."""
     try:
         parsed = int(value)
-    except (TypeError, ValueError) as exc:
-        raise ValueError(f"Field '{field_name}' must be a positive integer.") from exc
+    except (TypeError, ValueError):
+        return f"Field '{field_name}' must be a positive integer."
     if parsed <= 0:
-        raise ValueError(f"Field '{field_name}' must be a positive integer.")
-    return parsed
+        return f"Field '{field_name}' must be a positive integer."
+    return None
 
 
-def coerce_positive_int(value, default: int) -> int:
-    """Permissive positive-int coercion used for non-boundary defaults."""
+def parse_positive_int(value, default: int) -> int:
+    """Parse value as positive int, returning default on failure."""
     try:
         parsed = int(value)
     except (TypeError, ValueError):
@@ -88,13 +45,10 @@ def coerce_positive_int(value, default: int) -> int:
     return max(1, parsed)
 
 
-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
+def parse_positive_float(value, default: float) -> float:
+    """Parse value as positive float, returning default on failure."""
     try:
         parsed = float(value)
     except (TypeError, ValueError):
-        return fallback
-    return parsed if parsed > 0.0 else fallback
+        return max(0.0, float(default))
+    return max(0.0, parsed)

api/serverless_routes.py

@@ -0,0 +1,98 @@
+"""Master-only Runpod pool monitoring. Health checks never submit inference."""
+from aiohttp import web
+import server
+
+from ..utils.config import load_config
+from .orchestration.serverless_runtime import get_manager
+from .orchestration.runpod_client import RunpodError
+
+
+def _pool(worker_id):
+    return next((w for w in load_config().get("workers", [])
+                 if w.get("id") == worker_id and w.get("type") == "serverless"), None)
+
+
+async def _status(worker_id, fresh=False):
+    worker = _pool(worker_id)
+    if worker is None:
+        return web.json_response({"ok": False, "message": "Serverless worker not found"}, status=404)
+    jobs = []
+    try:
+        manager = get_manager()
+        jobs = manager.public_jobs(worker_id)
+        result = await manager.health(worker, fresh=fresh)
+        return web.json_response(dict(result, jobs=jobs))
+    except Exception as exc:
+        message = str(exc) if isinstance(exc, RunpodError) else "Check the saved endpoint ID and API key"
+        return web.json_response({"ok": False, "queueRemaining": 0, "message": message, "jobs": jobs})
+
+
+@server.PromptServer.instance.routes.post("/distributed/serverless/test")
+async def test_serverless_endpoint(request):
+    try:
+        data = await request.json()
+        worker_id = data.get("worker_id")
+    except Exception:
+        return web.json_response({"ok": False, "message": "Invalid request"}, status=400)
+    return await _status(worker_id, fresh=True)
+
+
+@server.PromptServer.instance.routes.get("/distributed/serverless/status/{worker_id}")
+async def serverless_status_endpoint(request):
+    return await _status(request.match_info["worker_id"])
+
+
+@server.PromptServer.instance.routes.post("/distributed/serverless/cancel")
+async def cancel_serverless_endpoint(request):
+    try:
+        data = await request.json()
+        worker_id = data.get("worker_id")
+    except Exception:
+        return web.json_response({"ok": False, "message": "Invalid request"}, status=400)
+    if _pool(worker_id) is None:
+        return web.json_response({"ok": False, "message": "Serverless worker not found"}, status=404)
+    try:
+        manager = get_manager()
+        count = await manager.cancel(worker_id=worker_id)
+        return web.json_response({"ok": True, "cancelled": count, "jobs": manager.public_jobs(worker_id),
+                                  "message": "Cancellation requested for this master's jobs only"})
+    except Exception:
+        return web.json_response({"ok": False, "message": "Cancellation could not be confirmed; check provider jobs"}, status=502)
+
+
+async def _recover(request, *, acknowledge=False):
+    try:
+        data = await request.json()
+        worker_id = data.get("worker_id")
+        if not isinstance(worker_id, str) or not worker_id.strip():
+            raise ValueError()
+        record_id = data.get("record_id")
+        if acknowledge:
+            if not isinstance(record_id, str) or not record_id.strip() or data.get("provider_checked") is not True:
+                raise ValueError()
+    except Exception:
+        return web.json_response({"ok": False, "message": "Provide a worker ID; manual resolution also requires a record ID and provider_checked: true"}, status=400)
+    if _pool(worker_id) is None:
+        return web.json_response({"ok": False, "message": "Serverless worker not found"}, status=404)
+    try:
+        manager = get_manager()
+        if acknowledge:
+            await manager.acknowledge_recovery(worker_id, record_id, provider_checked=True)
+            result = {"message": "Recovery acknowledged by operator; not provider-confirmed cancellation"}
+        else:
+            result = {"resolved": await manager.reconcile(worker_id=worker_id)}
+        return web.json_response(dict(result, ok=True, jobs=manager.public_jobs(worker_id)))
+    except ValueError:
+        return web.json_response({"ok": False, "message": "Unresolved inactive job not found for this worker"}, status=409)
+    except Exception:
+        return web.json_response({"ok": False, "message": "Recovery could not be confirmed; retained jobs remain protected"}, status=502)
+
+
+@server.PromptServer.instance.routes.post("/distributed/serverless/reconcile")
+async def reconcile_serverless_endpoint(request):
+    return await _recover(request)
+
+
+@server.PromptServer.instance.routes.post("/distributed/serverless/resolve")
+async def resolve_serverless_endpoint(request):
+    return await _recover(request, acknowledge=True)

api/tunnel_routes.py

@@ -1,63 +1,51 @@
-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:
+async def tunnel_status_endpoint(request):
     """Return Cloudflare tunnel status and last known details."""
-
-    async def _operation() -> web.StreamResponse:
+    try:
         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,
-            }
-        )
-
-    return await run_authorized_endpoint(request, _operation)
-
+        return web.json_response({
+            "status": "success",
+            "tunnel": status,
+            "master_host": master_host
+        })
+    except Exception as e:
+        return await handle_api_error(request, e, 500)
 
 @server.PromptServer.instance.routes.post("/distributed/tunnel/start")
-async def tunnel_start_endpoint(request: web.Request) -> web.StreamResponse:
+async def tunnel_start_endpoint(request):
     """Start a Cloudflare tunnel pointing at the current ComfyUI server."""
-
-    async def _operation() -> web.StreamResponse:
+    try:
         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"),
-            }
-        )
-
-    return await run_authorized_endpoint(request, _operation)
-
+        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)
 
 @server.PromptServer.instance.routes.post("/distributed/tunnel/stop")
-async def tunnel_stop_endpoint(request: web.Request) -> web.StreamResponse:
+async def tunnel_stop_endpoint(request):
     """Stop the managed Cloudflare tunnel if running."""
-
-    async def _operation() -> web.StreamResponse:
+    try:
         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"),
-            }
-        )
-
-    return await run_authorized_endpoint(request, _operation)
+        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)

api/usdu_routes.py

@@ -1,5 +1,3 @@
-from __future__ import annotations
-
 import asyncio
 import io
 import time
@@ -8,78 +6,15 @@ 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 ensure_tile_jobs_initialized, init_dynamic_job
-from ..upscale.payload_parsers import parse_tiles_from_form
+from ..upscale.job_store import MAX_PAYLOAD_SIZE, ensure_tile_jobs_initialized
+from ..upscale.payload_parsers import _parse_tiles_from_form
 from ..utils.logging import debug_log
 from ..utils.network import handle_api_error
-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/init_dynamic_job")
-async def init_dynamic_job_endpoint(request: web.Request) -> web.StreamResponse:
-    """Allow workers to initialize a dynamic job queue on the master.
-
-    Called by the first worker to reach the USDU node in delegate-only mode.
-    Idempotent — subsequent calls are no-ops if the queue already exists.
-    """
-    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)
-
-        batch_size = _parse_int_field(data.get("batch_size"), "batch_size", minimum=1)
-        enabled_workers = data.get("enabled_workers") or []
-        if isinstance(enabled_workers, str):
-            import json as _json
-            enabled_workers = _json.loads(enabled_workers)
-
-        await init_dynamic_job(multi_job_id, batch_size, enabled_workers)
-        return web.json_response({"status": "success"})
-    except Exception as e:
-        return await handle_api_error(request, e, 500)
 
 
 @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
+async def heartbeat_endpoint(request):
     try:
         data = await request.json()
         worker_id = data.get('worker_id')
@@ -93,8 +28,6 @@ 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"})
@@ -105,38 +38,24 @@ 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:
+async def submit_tiles_endpoint(request):
     """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:
-            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)
+        if content_length and int(content_length) > MAX_PAYLOAD_SIZE:
+            return await handle_api_error(request, f"Payload too large: {content_length} bytes", 413)
 
         data = await request.post()
         multi_job_id = data.get('multi_job_id')
         worker_id = data.get('worker_id')
-        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)
+        is_last = data.get('is_last', 'False').lower() == 'true'
 
         if multi_job_id is None or worker_id is None:
             return await handle_api_error(request, "Missing multi_job_id or worker_id", 400)
 
         prompt_server = ensure_tile_jobs_initialized()
 
-        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)
+        batch_size = int(data.get('batch_size', 0))
 
         # Handle completion signal
         if batch_size == 0 and is_last:
@@ -145,8 +64,6 @@ 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,
@@ -156,7 +73,7 @@ async def submit_tiles_endpoint(request: web.Request) -> web.StreamResponse:
                     return web.json_response({"status": "success"})
 
         try:
-            tiles = parse_tiles_from_form(data)
+            tiles = _parse_tiles_from_form(data)
         except ValueError as e:
             return await handle_api_error(request, str(e), 400)
 
@@ -166,8 +83,6 @@ 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:
@@ -191,28 +106,17 @@ 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:
+async def submit_image_endpoint(request):
     """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:
-            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)
+        if content_length and int(content_length) > MAX_PAYLOAD_SIZE:
+            return await handle_api_error(request, f"Payload too large: {content_length} bytes", 413)
 
         data = await request.post()
         multi_job_id = data.get('multi_job_id')
         worker_id = data.get('worker_id')
-        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)
+        is_last = data.get('is_last', 'False').lower() == 'true'
 
         if multi_job_id is None or worker_id is None:
             return await handle_api_error(request, "Missing multi_job_id or worker_id", 400)
@@ -221,10 +125,7 @@ async def submit_image_endpoint(request: web.Request) -> web.StreamResponse:
 
         # Handle image submission
         if 'full_image' in data and 'image_idx' in data:
-            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)
+            image_idx = int(data.get('image_idx'))
             img_data = data['full_image'].file.read()
             img = Image.open(io.BytesIO(img_data)).convert("RGB")
 
@@ -235,8 +136,6 @@ 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,
@@ -252,8 +151,6 @@ 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,
@@ -269,11 +166,8 @@ 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:
+async def request_image_endpoint(request):
     """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')
@@ -296,8 +190,6 @@ 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)
@@ -307,33 +199,25 @@ 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(
-                            {
-                                "kind": "image",
-                                "task_idx": task_idx,
-                                "estimated_remaining": remaining,
-                            }
-                        )
+                        return web.json_response({"image_idx": task_idx, "estimated_remaining": remaining})
                     debug_log(f"UltimateSDUpscale API - Assigned tile {task_idx} to worker {worker_id}")
                     return web.json_response({
-                        "kind": "tile",
-                        "task_idx": task_idx,
+                        "tile_idx": task_idx,
                         "estimated_remaining": remaining,
                         "batched_static": job_data.batched_static,
                     })
                 except asyncio.TimeoutError:
-                    return web.json_response({"kind": "none", "task_idx": None, "estimated_remaining": 0})
+                    if mode == 'dynamic':
+                        return web.json_response({"image_idx": None})
+                    return web.json_response({"tile_idx": None})
             return await handle_api_error(request, "Job not found", 404)
     except Exception as e:
         return await handle_api_error(request, e, 500)
 
 
 @server.PromptServer.instance.routes.get("/distributed/job_status")
-async def job_status_endpoint(request: web.Request) -> web.StreamResponse:
+async def job_status_endpoint(request):
     """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

@@ -1,13 +1,10 @@
-from __future__ import annotations
-
 import json
 import asyncio
 import os
 import time
 import platform
-import subprocess  # nosec B404 - commands are fixed and never shell-expanded
+import subprocess
 import socket
-from typing import Any
 
 import torch
 import aiohttp
@@ -19,33 +16,35 @@ from ..utils.logging import debug_log, log
 from ..utils.network import (
     build_worker_url,
     get_client_session,
+    get_server_port,
     handle_api_error,
     normalize_host,
     probe_worker,
 )
 from ..utils.constants import CHUNK_SIZE
 from ..workers import get_worker_manager
-from .request_guards import authorization_error_or_none
-from .schemas import (
-    distributed_auth_headers,
-    require_fields,
-    require_worker_id,
-)
+from ..workers.ports import allocate_worker_ports
+from .schemas import require_fields, validate_worker_id
 from ..workers.detection import (
     get_machine_id,
     is_docker_environment,
     is_runpod_environment,
 )
-from ..utils.async_helpers import PromptValidationError, queue_prompt_payload
+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 {}
 
 
 @server.PromptServer.instance.routes.get("/distributed/worker_ws")
-async def worker_ws_endpoint(request: web.Request) -> web.StreamResponse:
+async def worker_ws_endpoint(request):
     """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)
 
@@ -116,11 +115,8 @@ async def worker_ws_endpoint(request: web.Request) -> web.StreamResponse:
 
 
 @server.PromptServer.instance.routes.post("/distributed/worker/clear_launching")
-async def clear_launching_state(request: web.Request) -> web.StreamResponse:
+async def clear_launching_state(request):
     """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()
@@ -130,10 +126,8 @@ async def clear_launching_state(request: web.Request) -> web.StreamResponse:
 
         worker_id = str(data.get("worker_id")).strip()
         config = load_config()
-        try:
-            worker_id = require_worker_id(worker_id, config)
-        except ValueError as exc:
-            return await handle_api_error(request, exc, 404)
+        if not validate_worker_id(worker_id, config):
+            return await handle_api_error(request, f"Worker {worker_id} not found", 404)
         
         # Clear launching flag in managed processes
         if worker_id in wm.processes:
@@ -147,9 +141,8 @@ async def clear_launching_state(request: web.Request) -> web.StreamResponse:
         return await handle_api_error(request, e, 500)
 
 
-def get_network_ips() -> list[str]:
+def get_network_ips():
     """Get all network IPs, trying multiple methods."""
-    command_timeout = 5.0
     ips = []
     hostname = socket.gethostname()
 
@@ -160,8 +153,8 @@ def get_network_ips() -> list[str]:
             ip = info[4][0]
             if ip and ip not in ips and not ip.startswith('::'):  # Skip IPv6 for now
                 ips.append(ip)
-    except (socket.gaierror, OSError) as exc:
-        debug_log(f"get_network_ips: getaddrinfo failed for hostname {hostname}: {exc}")
+    except (socket.gaierror, OSError):
+        pass
 
     # Method 2: Try to connect to external server and get local IP
     try:
@@ -171,19 +164,14 @@ def get_network_ips() -> list[str]:
         s.close()
         if local_ip not in ips:
             ips.append(local_ip)
-    except (OSError, socket.error) as exc:
-        debug_log(f"get_network_ips: UDP local IP probe failed: {exc}")
+    except (OSError, socket.error):
+        pass
 
     # Method 3: Platform-specific commands
     try:
         if platform.system() == "Windows":
             # Windows ipconfig
-            result = subprocess.run(  # nosec B603 - static command, no user input
-                ["ipconfig"],
-                capture_output=True,
-                text=True,
-                timeout=command_timeout,
-            )
+            result = subprocess.run(["ipconfig"], capture_output=True, text=True)
             lines = result.stdout.split('\n')
             for i, line in enumerate(lines):
                 if 'IPv4' in line and i + 1 < len(lines):
@@ -193,20 +181,10 @@ def get_network_ips() -> list[str]:
         else:
             # Unix/Linux/Mac ifconfig or ip addr
             try:
-                result = subprocess.run(  # nosec B603 - static command, no user input
-                    ["ip", "addr"],
-                    capture_output=True,
-                    text=True,
-                    timeout=command_timeout,
-                )
+                result = subprocess.run(["ip", "addr"], capture_output=True, text=True)
             except (FileNotFoundError, OSError):
                 try:
-                    result = subprocess.run(  # nosec B603 - static command, no user input
-                        ["ifconfig"],
-                        capture_output=True,
-                        text=True,
-                        timeout=command_timeout,
-                    )
+                    result = subprocess.run(["ifconfig"], capture_output=True, text=True)
                 except (FileNotFoundError, OSError):
                     result = None
 
@@ -217,13 +195,13 @@ def get_network_ips() -> list[str]:
                     ip = match.group(1)
                     if ip and ip not in ips:
                         ips.append(ip)
-    except (OSError, subprocess.SubprocessError) as exc:
-        debug_log(f"get_network_ips: platform command probe failed: {exc}")
+    except (OSError, subprocess.SubprocessError):
+        pass
 
     return ips
 
 
-def get_recommended_ip(ips: list[str]) -> str | None:
+def get_recommended_ip(ips):
     """Choose the best IP for master-worker communication."""
     # Priority order:
     # 1. Private network ranges (192.168.x.x, 10.x.x.x, 172.16-31.x.x)
@@ -258,7 +236,7 @@ def get_recommended_ip(ips: list[str]) -> str | None:
         return None
 
 
-def _get_cuda_info() -> tuple[int | None, int, int]:
+def _get_cuda_info():
     """Detect CUDA device index and total physical GPU count.
 
     Returns (cuda_device, cuda_device_count, physical_device_count).
@@ -274,7 +252,7 @@ def _get_cuda_info() -> tuple[int | None, int, int]:
             if visible_devices:
                 cuda_device = visible_devices[0]
                 try:
-                    result = subprocess.run(  # nosec B603 - static command, no user input
+                    result = subprocess.run(
                         ['nvidia-smi', '--query-gpu=name', '--format=csv,noheader'],
                         capture_output=True,
                         text=True,
@@ -298,9 +276,16 @@ def _get_cuda_info() -> tuple[int | None, int, int]:
         return None, 0, 0
 
 
-def _collect_network_info_sync() -> dict[str, Any]:
+def _collect_network_info_sync():
     """Collect network/cuda info in a worker thread to avoid blocking route handlers."""
     cuda_device, cuda_device_count, physical_device_count = _get_cuda_info()
+    device_count = physical_device_count if physical_device_count > 0 else cuda_device_count
+    master_port = get_server_port()
+    config = load_config()
+    worker_ports = []
+    if not config.get("settings", {}).get("has_auto_populated_workers") and not config.get("workers"):
+        worker_count = device_count - (1 if cuda_device is not None and 0 <= cuda_device < device_count else 0)
+        worker_ports = allocate_worker_ports(master_port, config.get("workers", []), worker_count)
     hostname = socket.gethostname()
     all_ips = get_network_ips()
     recommended_ip = get_recommended_ip(all_ips)
@@ -309,11 +294,13 @@ def _collect_network_info_sync() -> dict[str, Any]:
         "all_ips": all_ips,
         "recommended_ip": recommended_ip,
         "cuda_device": cuda_device,
-        "cuda_device_count": physical_device_count if physical_device_count > 0 else cuda_device_count,
+        "cuda_device_count": device_count,
+        "master_port": master_port,
+        "local_worker_ports": worker_ports,
     }
 
 
-def _read_worker_log_sync(log_file: str, lines_to_read: int) -> dict[str, Any]:
+def _read_worker_log_sync(log_file, lines_to_read):
     """Read worker log content from disk in a threadpool worker."""
     file_size = os.path.getsize(log_file)
 
@@ -349,12 +336,7 @@ def _read_worker_log_sync(log_file: str, lines_to_read: int) -> dict[str, Any]:
     }
 
 
-def _parse_positive_int_query(
-    value: Any,
-    default: int,
-    minimum: int = 1,
-    maximum: int | None = 10000,
-) -> int:
+def _parse_positive_int_query(value, default, minimum=1, maximum=10000):
     """Parse bounded positive integer query params with sane fallback."""
     try:
         parsed = int(value)
@@ -366,7 +348,7 @@ def _parse_positive_int_query(
     return parsed
 
 
-def _find_worker_by_id(config: dict[str, Any], worker_id: str) -> dict[str, Any] | None:
+def _find_worker_by_id(config, worker_id):
     worker_id_str = str(worker_id).strip()
     for worker in config.get("workers", []):
         if str(worker.get("id")).strip() == worker_id_str:
@@ -375,11 +357,8 @@ 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:
+async def get_local_log_endpoint(request):
     """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:
@@ -423,11 +402,8 @@ 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:
+async def get_network_info_endpoint(request):
     """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)
@@ -441,11 +417,8 @@ async def get_network_info_endpoint(request: web.Request) -> web.StreamResponse:
         return await handle_api_error(request, e, 500)
 
 @server.PromptServer.instance.routes.get("/distributed/system_info")
-async def get_system_info_endpoint(request: web.Request) -> web.StreamResponse:
+async def get_system_info_endpoint(request):
     """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
         
@@ -468,11 +441,8 @@ async def get_system_info_endpoint(request: web.Request) -> web.StreamResponse:
         return await handle_api_error(request, e, 500)
 
 @server.PromptServer.instance.routes.post("/distributed/launch_worker")
-async def launch_worker_endpoint(request: web.Request) -> web.StreamResponse:
+async def launch_worker_endpoint(request):
     """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()
@@ -484,14 +454,15 @@ async def launch_worker_endpoint(request: web.Request) -> web.StreamResponse:
         
         # Find worker config
         config = load_config()
-        try:
-            worker_id = require_worker_id(worker_id, config)
-        except ValueError as exc:
-            return await handle_api_error(request, exc, 404)
+        if not validate_worker_id(worker_id, config):
+            return await handle_api_error(request, f"Worker {worker_id} not found", 404)
         worker = next((w for w in config.get("workers", []) if str(w.get("id")) == worker_id), None)
         if not worker:
             return await handle_api_error(request, f"Worker {worker_id} not found", 404)
         
+        if worker.get("type") == "serverless":
+            return await handle_api_error(request, "Serverless workers are managed by Runpod, not local processes", 400)
+
         # Ensure consistent string ID
         worker_id_str = worker_id
         
@@ -506,7 +477,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)
@@ -526,6 +497,8 @@ async def launch_worker_endpoint(request: web.Request) -> web.StreamResponse:
                 "message": f"Worker {worker['name']} launched",
                 "log_file": log_file
             })
+        except ValueError as e:
+            return await handle_api_error(request, f"Failed to launch worker: {str(e)}", 400)
         except Exception as e:
             return await handle_api_error(request, f"Failed to launch worker: {str(e)}", 500)
             
@@ -534,11 +507,8 @@ 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:
+async def stop_worker_endpoint(request):
     """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()
@@ -548,10 +518,8 @@ async def stop_worker_endpoint(request: web.Request) -> web.StreamResponse:
 
         worker_id = str(data.get("worker_id")).strip()
         config = load_config()
-        try:
-            worker_id = require_worker_id(worker_id, config)
-        except ValueError as exc:
-            return await handle_api_error(request, exc, 404)
+        if not validate_worker_id(worker_id, config):
+            return await handle_api_error(request, f"Worker {worker_id} not found", 404)
 
         success, message = wm.stop_worker(worker_id)
         
@@ -569,11 +537,8 @@ 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:
+async def get_managed_workers_endpoint(request):
     """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({
@@ -585,16 +550,16 @@ 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:
+async def get_local_worker_status_endpoint(request):
     """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 = {}
         
         for worker in config.get("workers", []):
+            # Serverless pools have no host, port, or local process.
+            if worker.get("type") == "serverless":
+                continue
             # Only check local workers
             host = normalize_host(worker.get("host")) or ""
             if not host or host in ["localhost", "127.0.0.1"]:
@@ -658,11 +623,8 @@ 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:
+async def get_worker_log_endpoint(request):
     """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']
@@ -704,11 +666,8 @@ 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:
+async def get_remote_worker_log_endpoint(request):
     """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()
@@ -731,7 +690,6 @@ 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

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

bootstrap/entrypoint.py

@@ -1,66 +0,0 @@
-"""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,4 +1,28 @@
-"""Project-level pytest collection rules for plugin-style package layout."""
+# conftest.py — project-level pytest configuration.
+#
+# Problem: custom_nodes/ComfyUI-Distributed/__init__.py uses relative imports
+# (from .nodes.v3 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
 
 collect_ignore = [
     "__init__.py",

distributed.py

@@ -1,15 +0,0 @@
-"""Compatibility facade for legacy imports from `distributed`."""
-from __future__ import annotations
-
-from .bootstrap.entrypoint import (
-    build_node_mappings,
-    initialize_runtime,
-)
-
-NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS = build_node_mappings()
-
-__all__ = [
-    "NODE_CLASS_MAPPINGS",
-    "NODE_DISPLAY_NAME_MAPPINGS",
-    "initialize_runtime",
-]

docs/comfyui-distributed-api.md

@@ -19,8 +19,8 @@ This document describes the **public HTTP API** added to ComfyUI-Distributed to
 - `POST /distributed/queue` — queues a workflow using the same distributed orchestration rules as the UI:
   - Detects distributed nodes in the prompt (`DistributedCollector`, `UltimateSDUpscaleDistributed`).
   - Resolves enabled/selected workers.
-  - Pings workers (`GET /prompt`) to include only reachable ones.
-  - Dispatches the workflow to workers (`POST /prompt`).
+  - Probes and dispatches workers through `/distributed/worker_ws` by default.
+  - If `settings.websocket_orchestration=false`, probes with `GET /prompt` and dispatches with `POST /prompt` instead.
   - Queues the master workflow in ComfyUI’s prompt queue.
   - If any `DistributedCollector` has `load_balance=true`, selects one least-busy participant for this run.
 
@@ -53,6 +53,7 @@ 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"
 }
 ```
@@ -60,7 +61,8 @@ Queue a workflow for distributed execution.
 #### Fields
 
 - `prompt` (required unless `workflow.prompt` is present, object)
-  - The ComfyUI prompt/workflow graph, same shape as used by `POST /prompt`.
+  - A complete ComfyUI API-format prompt graph, using the same shape as `POST /prompt`.
+  - This is not the normal visual workflow export from the ComfyUI editor.
 - `workflow` (optional, object)
   - Workflow metadata that ComfyUI normally stores in `extra_pnginfo.workflow`.
   - If you don’t care about UI metadata, you can omit it.
@@ -73,6 +75,10 @@ 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>]`.
@@ -100,7 +106,8 @@ $cfg.workers | Select-Object id,name,enabled,host,port,type | Format-Table -Auto
 ```json
 {
   "prompt_id": "<uuid>",
-  "worker_count": 2
+  "worker_count": 2,
+  "auto_prepare_supported": true
 }
 ```
 
@@ -117,10 +124,14 @@ $cfg.workers | Select-Object id,name,enabled,host,port,type | Format-Table -Auto
 
 ## Worker requirements (important)
 
-For a worker to participate, it must be reachable from the master:
+For a worker to participate, it must be reachable from the master. By default:
+
+- WebSocket probe and dispatch: `<worker-base>/distributed/worker_ws` must accept the connection.
+
+If `settings.websocket_orchestration=false`:
 
 - Health check: `GET <worker-base>/prompt` must return HTTP 200.
-- Dispatch: `POST <worker-base>/prompt` must accept the workflow.
+- Dispatch: `POST <worker-base>/prompt` must accept the prompt.
 
 Also, for collector-based flows:
 
@@ -195,7 +206,7 @@ Proxy endpoint on master that fetches logs from a configured remote/cloud worker
 
 ## Examples
 
-### 1) Minimal `curl`
+### 1) Minimal `curl` request envelope
 
 ```bash
 curl -X POST "http://127.0.0.1:8188/distributed/queue" \
@@ -203,18 +214,23 @@ curl -X POST "http://127.0.0.1:8188/distributed/queue" \
   -d @payload.json
 ```
 
-Where `payload.json` contains at least:
+`payload.json` must contain a complete ComfyUI API-format prompt. The abbreviated envelope below illustrates the request shape but is not directly executable:
 
 ```json
 {
   "prompt": {
-    "1": {"class_type": "KSampler", "inputs": {} }
+    "<node_id>": {
+      "class_type": "<node class>",
+      "inputs": {"<required_input>": "<value or connection>"}
+    }
   },
   "enabled_worker_ids": [],
   "client_id": "external-client"
 }
 ```
 
+Export or construct a valid API-format prompt with all required node inputs and at least one output node before submitting it.
+
 ### 2) Python (`requests`)
 
 ```python

docs/compressed-result-transfer.md

@@ -0,0 +1,38 @@
+# Compressed Result Transfer
+
+**Distributed Collector** has two optional widgets that control how workers send their images back to the master. This matters most when workers run in the cloud and you pay for data leaving it (egress).
+
+| `transfer_format` | Type | Notes |
+|---|---|---|
+| `legacy_png` (default) | lossless | Original behaviour: one JSON request per frame, uncompressed PNG in base64. Works with any master version. |
+| `png` | lossless | Compressed PNG, binary, several frames per request. Same tensors as `legacy_png`. |
+| `ffv1` | lossless | Whole batch as one FFV1 video. Same tensors as `legacy_png`; best lossless choice for frame sequences. |
+| `jpeg`, `webp` | lossy | Per-frame stills. Good for image batches. |
+| `h264`, `h265`, `av1` | lossy | Whole batch as one video. By far the smallest for video frame batches. |
+
+`transfer_quality` (1–100, default 90) applies to the lossy formats. JPEG/WebP use it directly; video codecs map it to CRF (90 → CRF 13 for h264/h265, CRF 16 for av1; 100 → CRF 0).
+
+Measured on a 121-frame 1280×720 batch (one worker's result):
+
+| Format | Sent | vs legacy | PSNR |
+|---|---|---|---|
+| `legacy_png` | 446 MB | 1× | lossless |
+| `png` | 163 MB | 2.7× | lossless |
+| `ffv1` | 128 MB | 3.5× | lossless |
+| `jpeg` q90 | 27 MB | 17× | 42.7 dB |
+| `webp` q90 | 13 MB | 33× | 40.6 dB |
+| `h264` q90 | 2.4 MB | 186× | 40.2 dB |
+| `h265` q90 | 1.0 MB | 434× | 40.1 dB |
+| `av1` q90 | 0.4 MB | 1031× | 40.5 dB |
+
+The test clip is a slow pan, which video codecs compress especially well; expect smaller (still large) gains on content with more motion.
+
+Notes:
+- Choose the format on the master; the value is sent to workers with the workflow.
+- The master must run this version to accept compressed results. A worker sending to an older master gets a 400 on its first request and automatically resends with `legacy_png`.
+- If a worker's PyAV build lacks the requested video encoder, it falls back to `h264` (lossy formats) or `png` (`ffv1`), and logs it.
+- Lossy formats only affect frames produced by workers; the master's own frames are never re-encoded.
+- Request size is bounded on the final encoded bytes. Frames are batched into requests of up to 32 MiB (`COMFYUI_DISTRIBUTED_TRANSFER_CHUNK_BYTES`); a single large frame may go alone in a request of up to 96 MiB (`COMFYUI_DISTRIBUTED_TRANSFER_MAX_REQUEST_BYTES`). Keep the latter below the master's `--max-upload-size` (default 100 MB) if you change either.
+- All chunks are encoded before anything is sent. If a single frame cannot fit within the request limit, the worker fails with an error and sends nothing; use a lossy format or raise both limits.
+- A rejection other than the older-master case above (for example the master's size limit, HTTP 413) is reported as an error rather than retried as `legacy_png`.
+- Each worker logs what it sent, e.g. `Sent 121 image(s) as h264 in 1 request(s): 2.40 MB (raw 334.49 MB, 139.4x smaller)`.

docs/serverless-workers.md

@@ -0,0 +1,173 @@
+# Runpod serverless workers (preview)
+
+A **Serverless (Runpod)** worker is a queue-based endpoint, not a permanent
+ComfyUI machine. Runpod chooses the actual GPU workers. Distributed submits one
+normal ComfyUI API workflow per requested image and collects the returned images.
+The endpoint workers **do not need ComfyUI-Distributed installed**.
+
+## Setup
+
+1. Add a worker in the Distributed panel and select **Serverless (Runpod)**.
+2. Enter its name, Runpod **endpoint ID** (not a URL; a pasted `https://api.runpod.ai/v2/<id>/…`
+   URL is trimmed to the ID), and API key.
+3. Start with **Images per run = 1** and **Timeout = 1800**.
+4. Save, then choose **Test Saved Endpoint**. This checks authentication and
+   endpoint health only; it does not start GPU inference or prove model/node compatibility.
+5. Enable the worker when you want queued Distributed workflows to use it.
+6. Use an image **Distributed Collector** and queue through Distributed as usual.
+   Enable **Delegate Only** if the master should collect without generating locally.
+
+Keys are write-only. Leave the key field blank to retain a saved key, or select
+**Clear saved API key** to remove it. The checkbox is unavailable while the
+worker is enabled, so turn the worker off first. A saved key is required before
+enabling a worker. Use HTTPS or a trusted local connection when
+entering keys. Like ComfyUI itself, the panel must not be exposed to untrusted
+users: anyone able to queue arbitrary workflows can run code and spend endpoint credits.
+
+## Endpoint contract
+
+The endpoint must accept this generic input shape; named application modes are
+not used:
+
+```json
+{
+  "input": {
+    "workflow": {
+      "1": {"class_type": "EmptyImage", "inputs": {"width": 32, "height": 32, "batch_size": 1, "color": 0}},
+      "2": {"class_type": "SaveImage", "inputs": {"images": ["1", 0], "filename_prefix": "example"}}
+    },
+    "images": []
+  }
+}
+```
+
+For image inputs, `images` contains `{ "name": "input.png", "image": "<base64>" }`
+entries and the handler must rewrite `LoadImage` references to the staged files.
+The adapter transfers local input images; it does not copy models or install nodes.
+All referenced models and custom nodes must already exist in the endpoint image.
+
+On completion, Runpod's `output.images` must contain artifact records such as:
+
+```json
+{
+  "filename": "example_00001_.png",
+  "subfolder": "",
+  "artifact_type": "output",
+  "artifact_path": "example_00001_.png",
+  "type": "base64",
+  "data": "<base64>"
+}
+```
+
+Image data can alternatively use `type: "s3_url"`, `"gcs_url"`, or `"url"`
+with an HTTPS URL in `data`/`url`. Signed artifact URLs are downloaded without
+Runpod credentials; private-network destinations and unsafe redirects are refused.
+The handler must preserve injected `SaveImage.filename_prefix` values so results
+can be matched to the right collector and slot. Duplicate basenames in different
+subfolders remain different artifact identities.
+
+A Runpod job marked `COMPLETED` is still rejected if `output.error` or
+`output.errors` is nonempty, required collector outputs are missing, or image
+artifacts are invalid. Failures are surfaced rather than silently returning an
+incomplete serverless batch.
+
+## Scheduling and cancellation
+
+- **Images per run** (1–16) is how many separate Runpod jobs this worker submits
+  for each queued run. Each returns one image and receives its own Distributed
+  Seed/Value overrides. All of them are submitted at once.
+- How many of those jobs actually run in parallel is decided by the Runpod
+  endpoint's **max workers**, not by Distributed. Jobs beyond that wait in
+  Runpod's queue. To run every image at the same time, keep Images per run at or
+  below the endpoint's max workers.
+- Distributed also keeps a fixed limit of 16 outstanding jobs per serverless
+  worker (and 64 across all of them). Retained jobs after a restart, ambiguous
+  submissions, and unconfirmed cancellations keep counting against it until
+  reconciled. Other applications using the same endpoint are unaffected.
+- Older configs may contain a `concurrency` field from the previous separate
+  setting. It is still accepted when loading or saving, but ignored and dropped.
+- **Load balance** selects one participant, so it sends one job instead of
+  Images per run.
+- **Timeout** includes Runpod queue/cold-start time, execution, and output
+  retrieval. Provider policy limits are submitted as well.
+- A zero-worker endpoint can be healthy: Runpod scales from zero on submission.
+- Interrupting a collector or using the panel cancellation action requests
+  cancellation only for jobs submitted by this master. It never purges the endpoint.
+- A network failure during submission is **not retried automatically**: the job
+  may already be billable. Check the Runpod job list before queuing again.
+
+## Persistence and recovery
+
+Credentials live outside `gpu_config.json` and Git in the master ComfyUI user
+folder under `distributed/serverless/credentials.json`, with restrictive POSIX
+permissions. They are omitted from public configuration, workflow exports, and
+browser configuration state. The first version requires a POSIX (Linux/macOS)
+master for serverless secret storage; existing non-serverless workers remain
+available on Windows.
+
+`distributed/serverless/jobs.json` records provider job IDs and sanitized states,
+not image data, workflow payloads, or API keys. A restarted master cannot resume
+an interrupted collector. It does **not** automatically resubmit jobs. Retained
+IDs are available through `/distributed/serverless/status/<worker_id>`; the worker's
+cancel action can request cancellation of known pre-restart jobs. If no job ID
+was received, check Runpod manually. Cancellation is best-effort and may race
+with completion; `CANCEL_REQUESTED` is not a provider-confirmed cancellation.
+
+Unresolved records protect the worker's endpoint, API key, worker type,
+and deletion even when no local task remains. Disabling the worker is still allowed;
+it does not cancel existing jobs. On restart, never-submitted `WAITING` slots are
+closed locally, while interrupted `SUBMITTING` records without IDs become
+`UNKNOWN` and remain protected.
+
+### Reconcile or acknowledge recovery
+
+Recovery uses the same trusted-master API as the panel; it never resubmits a job.
+
+1. Read `GET /distributed/serverless/status/<worker_id>`. The `jobs` list includes
+   every unresolved record with its local `record_id`, provider `job_id` (if
+   received), and state. These records remain visible if endpoint health fails.
+2. Request cancellation from the panel, or send
+   `POST /distributed/serverless/reconcile` with the body below to **read** the
+   retained provider IDs. Only a confirmed terminal state releases protection.
+   Network errors, expired IDs, or ambiguous responses do not.
+
+   ```json
+   {"worker_id": "YOUR_WORKER_ID"}
+   ```
+
+3. If no ID was returned, or Runpod no longer retains its status, inspect Runpod
+   manually first. **Only after verifying that the job is no longer running or
+   queued**, send `POST /distributed/serverless/resolve` for that one record:
+
+   ```json
+   {"worker_id": "YOUR_WORKER_ID", "record_id": "LOCAL_RECORD_ID", "provider_checked": true}
+   ```
+
+   This records `RESOLVED`: an explicit operator acknowledgment, **not** a
+   provider-confirmed cancellation. It cannot resolve a live local task and does
+   not stop remote compute. It releases the record's reservation and may let
+   already-waiting jobs proceed. Never automate this acknowledgment as a retry
+   mechanism; an unverified job can still be billable.
+
+Worker setting changes become available once all its local tasks and unresolved records
+are finished. No credentials or ledger files need to be edited by hand.
+
+## First-version limits
+
+- Image collectors only. Audio/video transfer and audio collectors are rejected.
+- Coordinated `UltimateSDUpscaleDistributed` workloads and unsupported Distributed
+  node types are rejected before submission.
+- Downstream post-collector processing remains local; the remote graph contains
+  only the upstream compute required for image saves.
+- Inputs must be supported images under ComfyUI's input/output/temp roots.
+- Raw image input reads are guarded at 10,000,000 bytes total, and the complete
+  serialized `/run` JSON request body (workflow, base64 images, and policy) must
+  also fit within 10,000,000 bytes. Base64 adds about 33%, so the usable raw-image
+  budget is lower and depends on workflow size. Large-media/object-storage upload
+  support is not included.
+- No remote GPU discovery, endpoint provisioning, model synchronization, automatic
+  scaling changes, or inbound master callback/tunnel is needed.
+
+This feature is verified with offline provider/transport tests. A successful
+health test is not a paid end-to-end generation test; test your deployed handler
+and model inventory with one image per run before increasing it.

docs/v3-migration.md

@@ -0,0 +1,53 @@
+# Native V3 node API
+
+The package registers one `ComfyExtension` through `comfy_entrypoint()` and
+imports the versioned `comfy_api.v0_0_2` API. Use a ComfyUI version that provides
+that API; there is no V1 registration fallback.
+
+## Compatibility
+
+- All eight node IDs, display names, categories, visible input order, defaults
+  and output order are retained. Existing execution algorithms remain in the
+  private collector, utilities and upscale modules.
+- The image/audio dividers explicitly declare ten typed outputs. Their existing
+  frontend extensions still show only the selected number of outputs. This
+  replaces the V1 `ByPassTypeTuple` indexing workaround without changing saved
+  IMAGE/AUDIO socket indices or the ten returned values.
+- Standard hidden context uses `cls.hidden`. Worker/orchestration metadata keeps
+  its existing prompt-input names and Python defaults via `accept_all_inputs`;
+  it does not become visible widgets.
+- The collector retains list-input handling. Upscale retains its always-changing
+  fingerprint and creates a private runtime helper for each V3 execution, so
+  mutable helper state is not attached to sanitized V3 class clones.
+- Routes, distributed state and the existing worker startup/shutdown hooks are
+  initialized by `runtime/bootstrap.py` during `ComfyExtension.on_load()`.
+  Worker mode still suppresses automatic worker launch. `distributed.py` is no
+  longer an entrypoint.
+
+## Verification
+
+Run ordinary unit tests from this repository:
+
+```bash
+python -m pytest tests -q -o addopts=
+```
+
+Opt into real-framework acceptance using a ComfyUI checkout and its interpreter:
+
+```bash
+COMFYUI_SOURCE_ROOT=/path/to/ComfyUI \
+  /path/to/ComfyUI/.venv/bin/python -m pytest tests -q -o addopts=
+```
+
+The acceptance subprocess uses CPU mode and the real ComfyUI loader, input
+parser, V3 class preparation, prompt validation and `PromptExecutor`. It compares
+all node schemas with the V1 fixture, checks all five bundled workflows' node
+IDs and socket/link contracts, exercises injected worker metadata, image/audio
+lists, collector aggregation and divider outputs, rejects invalid upscale enums,
+and decodes an actual preview PNG referenced by executor history.
+
+The upscale GPU/model boundary is mocked to check argument forwarding and
+per-execution helper isolation. Full checkpoint inference, browser canvas
+acceptance and multi-host HTTP transport are not exercised. No HTTP listener,
+workers or model downloads are started; preview files use a temporary scratch
+directory. The test does not install this branch into a live custom-node folder.

docs/worker-setup-guides.md

@@ -74,12 +74,14 @@ The master can either contribute GPU work or stay in **orchestrator-only** mode:
 
 > ComfyUI instances running on a cloud service like Runpod. 
 
+For faster result transfers and lower bandwidth usage, see [Compressed Result Transfer](compressed-result-transfer.md).
+
 ### Deploy Cloud Worker on Runpod
 
 📺 [Watch Tutorial](https://www.youtube.com/watch?v=wxKKWMQhYTk)
 
 **On Runpod:**
-> If using your own template, make sure you launch ComfyUI with the `--enable-cors-header` argument and you `git clone ComfyUI-Distributed` into custom_nodes. ⚠️ **Required!**
+> If using your own template, launch ComfyUI with `--listen --enable-cors-header` and clone `ComfyUI-Distributed` into `custom_nodes`. ⚠️ **Required!**
 
 1. Register a [Runpod](https://get.runpod.io/0bw29uf3ug0p) account.
 2. On Runpod, go to Storage > New Network Volume and create a volume that will store the models you need. Start with 40 GB, you can always add more later. Learn more [about Network Volumes](https://docs.runpod.io/pods/storage/create-network-volumes).
@@ -92,7 +94,7 @@ The master can either contribute GPU work or stay in **orchestrator-only** mode:
 	- SAGE_ATTENTION: optional optimisation (set to true/false)
 5. Deploy your pod.
 6. Connect to your pod using JupyterLabs. This gives us access to the pod's file system.
-7. Download models into /workspaces/ComfyUI/models/ (these will remain on your network drive even after you terminate the pod). Example commands below:
+7. Download models into `/workspace/ComfyUI/models/` (these will remain on your network drive even after you terminate the pod). Example commands below:
 ```
 # Download from CivitAI
 comfy model download --url https://civitai.com/api/download/models/1759168 --relative-path /workspace/ComfyUI/models/checkpoints --set-civitai-api-token $CIVITAI_API_TOKEN

nodes/__init__.py

@@ -1,37 +1 @@
-from .utilities import (
-    DistributedSeed,
-    DistributedModelName,
-    DistributedValue,
-    ImageBatchDivider,
-    AudioBatchDivider,
-    DistributedEmptyImage,
-    AnyType,
-    ByPassTypeTuple,
-    any_type,
-)
-from .collector import DistributedCollectorNode
-from .branch import DistributedBranch
-from .branch_collector import DistributedBranchCollector
-
-NODE_CLASS_MAPPINGS = {
-    "DistributedCollector": DistributedCollectorNode,
-    "DistributedBranch": DistributedBranch,
-    "DistributedBranchCollector": DistributedBranchCollector,
-    "DistributedSeed": DistributedSeed,
-    "DistributedModelName": DistributedModelName,
-    "DistributedValue": DistributedValue,
-    "ImageBatchDivider": ImageBatchDivider,
-    "AudioBatchDivider": AudioBatchDivider,
-    "DistributedEmptyImage": DistributedEmptyImage,
-}
-NODE_DISPLAY_NAME_MAPPINGS = {
-    "DistributedCollector": "Distributed Collector",
-    "DistributedBranch": "Distributed Branch",
-    "DistributedBranchCollector": "Distributed Branch Collector",
-    "DistributedSeed": "Distributed Seed",
-    "DistributedModelName": "Distributed Model Name",
-    "DistributedValue": "Distributed Value",
-    "ImageBatchDivider": "Image Batch Divider",
-    "AudioBatchDivider": "Audio Batch Divider",
-    "DistributedEmptyImage": "Distributed Empty Image",
-}
+"""Private execution helpers; public registration lives in nodes.v3."""

nodes/branch.py

@@ -1,33 +0,0 @@
-from typing import Any
-
-from .utilities import any_type
-
-
-class DistributedBranch:
-    @classmethod
-    def INPUT_TYPES(cls: type["DistributedBranch"]) -> dict[str, Any]:
-        return {
-            "required": {
-                "input": (any_type,),
-                "num_branches": ("INT", {"default": 2, "min": 2, "max": 10, "step": 1}),
-            },
-            "hidden": {
-                "is_worker": ("BOOLEAN", {"default": False}),
-                "worker_id": ("STRING", {"default": ""}),
-                "assigned_branch": ("INT", {"default": -1, "min": -1, "max": 9}),
-                "multi_job_id": ("STRING", {"default": ""}),
-            },
-        }
-
-    RETURN_TYPES = tuple([any_type] * 10)
-    RETURN_NAMES = tuple([f"branch_{idx + 1}" for idx in range(10)])
-    FUNCTION = "branch"
-    CATEGORY = "utils"
-
-    def branch(
-        self,
-        input: Any,
-        num_branches: int = 2,
-        **_kwargs: Any,
-    ) -> tuple[Any, ...]:
-        return tuple([input] * 10)

nodes/branch_collector.py

@@ -1,298 +0,0 @@
-import asyncio
-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, 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 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
-from .runtime_helpers import (
-    get_prompt_server_instance as _get_prompt_server_instance,
-    throw_if_processing_interrupted as _throw_if_processing_interrupted,
-)
-from .utilities import any_type
-
-
-MAX_BRANCH_OUTPUTS = 10
-
-
-@dataclass(frozen=True)
-class BranchRunContext:
-    multi_job_id: str = ""
-    is_worker: bool = False
-    master_url: str = ""
-    enabled_worker_ids: list[str] = field(default_factory=list)
-    worker_id: str = ""
-    assigned_branch: int = -1
-    delegate_only: bool = False
-
-class DistributedBranchCollector:
-    @classmethod
-    def INPUT_TYPES(cls: type["DistributedBranchCollector"]) -> dict[str, Any]:
-        optional_inputs = {
-            f"branch_{idx + 1}": (any_type,)
-            for idx in range(MAX_BRANCH_OUTPUTS)
-        }
-        return {
-            "required": {
-                "num_branches": ("INT", {"default": 2, "min": 2, "max": MAX_BRANCH_OUTPUTS, "step": 1}),
-            },
-            "optional": optional_inputs,
-            "hidden": build_distributed_hidden_inputs(
-                include_assigned_branch=True,
-                assigned_branch_max=MAX_BRANCH_OUTPUTS - 1,
-            ),
-        }
-
-    RETURN_TYPES = tuple([any_type] * MAX_BRANCH_OUTPUTS)
-    RETURN_NAMES = tuple([f"branch_{idx + 1}" for idx in range(MAX_BRANCH_OUTPUTS)])
-    FUNCTION = "run"
-    CATEGORY = "utils"
-
-    def _branch_values(self, **kwargs):
-        values = []
-        for idx in range(MAX_BRANCH_OUTPUTS):
-            values.append(kwargs.get(f"branch_{idx + 1}", None))
-        return values
-
-    def _build_run_context(self, **kwargs: Any) -> BranchRunContext:
-        try:
-            assigned_branch = int(kwargs.get("assigned_branch", -1))
-        except (TypeError, ValueError):
-            assigned_branch = -1
-
-        common_context = parse_distributed_hidden_context(kwargs)
-        return BranchRunContext(
-            assigned_branch=assigned_branch,
-            **common_context,
-        )
-
-    def _resolve_local_branch_value(self, branch_values, assigned_branch):
-        try:
-            assigned_idx = int(assigned_branch)
-        except (TypeError, ValueError):
-            assigned_idx = -1
-
-        if 0 <= assigned_idx < MAX_BRANCH_OUTPUTS:
-            assigned_value = branch_values[assigned_idx]
-            if assigned_value is not None:
-                return assigned_value, assigned_idx
-
-            # In participant-pruned prompts, each participant should typically have
-            # exactly one local branch input. If socket names drifted (e.g. branch_2
-            # connected for a participant assigned to slot 0), remap that sole value
-            # to the assigned slot so the participant still contributes correctly.
-            non_none_indices = [idx for idx, value in enumerate(branch_values) if value is not None]
-            if len(non_none_indices) == 1:
-                return branch_values[non_none_indices[0]], assigned_idx
-
-        for idx, value in enumerate(branch_values):
-            if value is not None:
-                return value, idx
-        return None, assigned_idx
-
-    def _build_outputs(self, branch_values, num_branches):
-        outputs = [None] * MAX_BRANCH_OUTPUTS
-        try:
-            branch_count = int(num_branches)
-        except (TypeError, ValueError):
-            branch_count = 2
-        branch_count = max(2, min(branch_count, MAX_BRANCH_OUTPUTS))
-
-        for idx in range(branch_count):
-            value = branch_values[idx]
-            if value is not None:
-                outputs[idx] = value
-        return tuple(outputs)
-
-    def run(
-        self,
-        num_branches: int = 2,
-        **kwargs: Any,
-    ) -> tuple[Any, ...]:
-        branch_values = self._branch_values(**kwargs)
-        context = self._build_run_context(**kwargs)
-
-        if not context.multi_job_id:
-            return self._build_outputs(branch_values, num_branches)
-
-        return run_async_in_server_loop(
-            self.execute(
-                branch_values,
-                num_branches=num_branches,
-                context=context,
-            )
-        )
-
-    def _participant_branch_slots(self, enabled_workers, delegate_mode, num_branches):
-        participants = list(enabled_workers) if delegate_mode else ["master"] + list(enabled_workers)
-        participant_count = len(participants)
-        if participant_count <= 0:
-            return {}
-
-        slots_by_participant = {}
-        for branch_slot in range(int(num_branches)):
-            participant_id = str(participants[branch_slot % participant_count])
-            slots_by_participant.setdefault(participant_id, []).append(branch_slot)
-        return slots_by_participant
-
-    def _expected_worker_ids_for_branches(self, enabled_workers, delegate_mode, num_branches):
-        slots_by_participant = self._participant_branch_slots(enabled_workers, delegate_mode, num_branches)
-        expected = []
-        for participant_id, slots in slots_by_participant.items():
-            if participant_id == "master" or not slots:
-                continue
-            expected.append(str(participant_id))
-        return expected
-
-    def _fallback_for_missing_branch(self, source_value):
-        if isinstance(source_value, torch.Tensor):
-            tensor = source_value
-            if tensor.ndim == 3:
-                tensor = tensor.unsqueeze(0)
-            if tensor.ndim == 4:
-                fallback = torch.zeros_like(tensor)
-                if fallback.is_cuda:
-                    fallback = fallback.cpu()
-                return ensure_contiguous(fallback)
-        return torch.zeros((1, 64, 64, 3), dtype=torch.float32)
-
-    async def _send_branch_result_to_master(self, value, branch_idx, multi_job_id, master_url, worker_id):
-        if not isinstance(value, torch.Tensor):
-            raise ValueError("DistributedBranchCollector currently supports torch.Tensor results for worker transfer.")
-
-        image_batch = value
-        if image_batch.ndim == 3:
-            image_batch = image_batch.unsqueeze(0)
-        if image_batch.ndim != 4 or image_batch.shape[0] <= 0:
-            raise ValueError(
-                "DistributedBranchCollector tensor result must be IMAGE-like with shape [B,H,W,C] or [H,W,C]."
-            )
-
-        payload = {
-            "job_id": str(multi_job_id),
-            "worker_id": str(worker_id),
-            "batch_idx": int(branch_idx),
-            "image": encode_tensor_png_data_url(image_batch, 0),
-            "is_last": True,
-        }
-
-        session = await get_client_session()
-        url = f"{master_url}/distributed/job_complete"
-        async with session.post(
-            url,
-            json=payload,
-            headers=distributed_auth_headers(load_config()),
-            timeout=aiohttp.ClientTimeout(total=60),
-        ) as response:
-            response.raise_for_status()
-
-    async def execute(self, *args: Any, **kwargs: Any) -> tuple[Any, ...]:
-        """Compatibility wrapper with a uniform execute() signature across collectors."""
-        return await self._execute_branch(*args, **kwargs)
-
-    async def _execute_branch(
-        self,
-        branch_values: list[Any],
-        num_branches: int = 2,
-        context: BranchRunContext | None = None,
-    ) -> tuple[Any, ...]:
-        run_context = context or BranchRunContext()
-        try:
-            branch_count = int(num_branches)
-        except (TypeError, ValueError):
-            branch_count = 2
-        branch_count = max(2, min(branch_count, MAX_BRANCH_OUTPUTS))
-
-        outputs = [None] * MAX_BRANCH_OUTPUTS
-        for idx in range(branch_count):
-            if branch_values[idx] is not None:
-                outputs[idx] = branch_values[idx]
-
-        local_value, local_branch_idx = self._resolve_local_branch_value(branch_values, run_context.assigned_branch)
-        if 0 <= local_branch_idx < MAX_BRANCH_OUTPUTS and local_value is not None:
-            outputs[local_branch_idx] = local_value
-
-        if run_context.is_worker:
-            if 0 <= local_branch_idx < MAX_BRANCH_OUTPUTS and local_value is not None:
-                try:
-                    await self._send_branch_result_to_master(
-                        local_value,
-                        local_branch_idx,
-                        run_context.multi_job_id,
-                        run_context.master_url,
-                        run_context.worker_id,
-                    )
-                    outputs[local_branch_idx] = local_value
-                except Exception as exc:
-                    log(f"Worker - DistributedBranchCollector failed to send branch result: {exc}")
-            return tuple(outputs)
-
-        delegate_mode = bool(run_context.delegate_only or is_master_delegate_only())
-        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)
-
-        if not expected_workers:
-            return tuple(outputs)
-
-        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()
-
-        workers_done: set[str] = set()
-
-        def _handle_queue_result(result: dict[str, Any]) -> None:
-            image_index = result.get("image_index")
-            tensor = result.get("tensor")
-            try:
-                slot_idx = int(image_index)
-            except (TypeError, ValueError):
-                slot_idx = -1
-
-            if 0 <= slot_idx < MAX_BRANCH_OUTPUTS and tensor is not None:
-                if isinstance(tensor, torch.Tensor):
-                    if tensor.is_cuda:
-                        tensor = tensor.cpu()
-                    tensor = ensure_contiguous(tensor)
-                outputs[slot_idx] = tensor
-
-        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 - DistributedBranchCollector 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)
-                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:
-            fallback = self._fallback_for_missing_branch(local_value)
-            for missing_worker_id in missing_workers:
-                for slot_idx in slots_by_participant.get(str(missing_worker_id), []):
-                    if 0 <= slot_idx < MAX_BRANCH_OUTPUTS and outputs[slot_idx] is None:
-                        outputs[slot_idx] = fallback
-
-        return tuple(outputs)

nodes/context_kwargs.py

@@ -1,26 +0,0 @@
-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", "") 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,28 +1,25 @@
+import json
 import math
 from functools import wraps
-from typing import Any, Callable
 
 import comfy.samplers
 
 from ..utils.logging import debug_log, log
 from ..utils.async_helpers import run_async_in_server_loop
 from ..upscale.job_store import ensure_tile_jobs_initialized
-from .hidden_inputs import build_distributed_hidden_inputs
 
 from ..upscale.tile_ops import TileOpsMixin
 from ..upscale.result_collector import ResultCollectorMixin
 from ..upscale.worker_comms import WorkerCommsMixin
 from ..upscale.job_state import JobStateMixin
-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]:
+def sync_wrapper(async_func):
     """Decorator to wrap async methods for synchronous execution."""
     @wraps(async_func)
-    def sync_func(self, *args: Any, **kwargs: Any) -> Any:
+    def sync_func(self, *args, **kwargs):
         # Use run_async_in_server_loop for ComfyUI compatibility
         return run_async_in_server_loop(
             async_func(self, *args, **kwargs),
@@ -30,10 +27,21 @@ def sync_wrapper(async_func: Callable[..., Any]) -> Callable[..., Any]:
         )
     return sync_func
 
-
-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)
+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 []
 
 class UltimateSDUpscaleDistributed(
     DynamicModeMixin,
@@ -48,22 +56,22 @@ class UltimateSDUpscaleDistributed(
     """
     Distributed version of Ultimate SD Upscale (No Upscale).
     
-    Supports three processing modes:
+    Supports two currently selected processing modes:
     1. Single GPU: No workers available, process everything locally
-    2. Static Mode: Small batches, distributes tiles across workers (flattened)
-    3. Dynamic Mode: Large batches, assigns whole images to workers dynamically
+    2. Distributed tile queue: Workers pull tiles from a shared queue
     
     Features:
-    - Multi-mode batch handling for efficient video/image upscaling
+    - Tile-based batch handling for video/image upscaling
     - Tiled VAE support for memory efficiency
-    - Dynamic load balancing for large batches
+    - Shared work queue so faster workers can process more tiles
     - Backward compatible with single-image workflows
     
     Environment Variables:
     - COMFYUI_MAX_BATCH: Chunk size for tile sending (default 20)
     - COMFYUI_MAX_PAYLOAD_SIZE: Max API payload bytes (default 50MB)
     
-    Threshold: dynamic_threshold input controls mode switch (default 8)
+    The hidden dynamic_threshold input is retained for workflow compatibility but
+    does not affect the current mode-selection policy.
     """
 
     def __init__(self):
@@ -73,14 +81,7 @@ class UltimateSDUpscaleDistributed(
         debug_log("UltimateSDUpscaleDistributed - Node initialized")
 
     @classmethod
-    def INPUT_TYPES(cls: type["UltimateSDUpscaleDistributed"]) -> dict[str, Any]:
-        hidden_inputs = build_distributed_hidden_inputs()
-        hidden_inputs.update(
-            {
-                "tile_indices": ("STRING", {"default": ""}),  # Unused - kept for compatibility
-                "dynamic_threshold": ("INT", {"default": 8, "min": 1, "max": 64}),
-            }
-        )
+    def INPUT_TYPES(s):
         return {
             "required": {
                 "upscaled_image": ("IMAGE",),
@@ -101,7 +102,15 @@ class UltimateSDUpscaleDistributed(
                 "force_uniform_tiles": ("BOOLEAN", {"default": True}),
                 "tiled_decode": ("BOOLEAN", {"default": False}),
             },
-            "hidden": hidden_inputs,
+            "hidden": {
+                "multi_job_id": ("STRING", {"default": ""}),
+                "is_worker": ("BOOLEAN", {"default": False}),
+                "master_url": ("STRING", {"default": ""}),
+                "enabled_worker_ids": ("STRING", {"default": "[]"}),
+                "worker_id": ("STRING", {"default": ""}),
+                "tile_indices": ("STRING", {"default": ""}),  # Unused - kept for compatibility
+                "dynamic_threshold": ("INT", {"default": 8, "min": 1, "max": 64}),
+            },
         }
 
     RETURN_TYPES = ("IMAGE",)
@@ -113,47 +122,12 @@ class UltimateSDUpscaleDistributed(
         """Force re-execution."""
         return float("nan")  # Always re-execute
 
-    def run(
-        self,
-        upscaled_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,
-        force_uniform_tiles: bool,
-        tiled_decode: bool,
-        multi_job_id: str = "",
-        is_worker: bool = False,
-        master_url: str = "",
-        enabled_worker_ids: str = "[]",
-        worker_id: str = "",
-        tile_indices: str = "",
-        dynamic_threshold: int = 8,
-    ) -> tuple[Any, ...]:
+    def run(self, upscaled_image, model, positive, negative, vae, seed, steps, cfg, 
+            sampler_name, scheduler, denoise, tile_width, tile_height, padding, 
+            mask_blur, force_uniform_tiles, tiled_decode,
+            multi_job_id="", is_worker=False, master_url="", enabled_worker_ids="[]", 
+            worker_id="", tile_indices="", dynamic_threshold=8):
         """Entry point - runs SYNCHRONOUSLY like Ultimate SD Upscaler."""
-        core_args = UpscaleCoreArgs(
-            model=model,
-            positive=positive,
-            negative=negative,
-            vae=vae,
-            seed=seed,
-            steps=steps,
-            cfg=cfg,
-            sampler_name=sampler_name,
-            scheduler=scheduler,
-            denoise=denoise,
-            tiled_decode=tiled_decode,
-        )
         # Strict WAN/FLOW batching: error if batch is not 4n+1 (except allow 1)
         try:
             batch_size = int(getattr(upscaled_image, 'shape', [1])[0])
@@ -168,15 +142,9 @@ class UltimateSDUpscaleDistributed(
             )
         if not multi_job_id:
             # No distributed processing, run single GPU version
-            return self.process_single_gpu(
-                upscaled_image=upscaled_image,
-                core_args=core_args,
-                tile_width=tile_width,
-                tile_height=tile_height,
-                padding=padding,
-                mask_blur=mask_blur,
-                force_uniform_tiles=force_uniform_tiles,
-            )
+            return self.process_single_gpu(upscaled_image, model, positive, negative, vae,
+                                          seed, steps, cfg, sampler_name, scheduler, denoise,
+                                          tile_width, tile_height, padding, mask_blur, force_uniform_tiles, tiled_decode)
         
         if is_worker:
             # Worker mode: process tiles synchronously
@@ -187,64 +155,24 @@ class UltimateSDUpscaleDistributed(
                                       worker_id, enabled_worker_ids, dynamic_threshold)
         else:
             # Master mode: distribute and collect synchronously
-            return self.process_master(
-                upscaled_image=upscaled_image,
-                core_args=core_args,
-                tile_width=tile_width,
-                tile_height=tile_height,
-                padding=padding,
-                mask_blur=mask_blur,
-                force_uniform_tiles=force_uniform_tiles,
-                multi_job_id=multi_job_id,
-                enabled_worker_ids=enabled_worker_ids,
-                dynamic_threshold=dynamic_threshold,
-            )
+            return self.process_master(upscaled_image, model, positive, negative, vae,
+                                     seed, steps, cfg, sampler_name, scheduler, denoise,
+                                     tile_width, tile_height, padding, mask_blur,
+                                     force_uniform_tiles, tiled_decode, multi_job_id, enabled_worker_ids, 
+                                     dynamic_threshold)
 
-    def process_worker(
-        self,
-        upscaled_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,
-        force_uniform_tiles: bool,
-        tiled_decode: bool,
-        multi_job_id: str,
-        master_url: str,
-        worker_id: str,
-        enabled_worker_ids: str,
-        dynamic_threshold: int,
-    ) -> tuple[Any, ...]:
+    def process_worker(self, upscaled_image, model, positive, negative, vae,
+                      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_worker_ids, dynamic_threshold):
         """Unified worker processing - handles both static and dynamic modes."""
-        core_args = UpscaleCoreArgs(
-            model=model,
-            positive=positive,
-            negative=negative,
-            vae=vae,
-            seed=seed,
-            steps=steps,
-            cfg=cfg,
-            sampler_name=sampler_name,
-            scheduler=scheduler,
-            denoise=denoise,
-            tiled_decode=tiled_decode,
-        )
         # Get batch size to determine mode
         batch_size = upscaled_image.shape[0]
         
         # Ensure mode consistency across master/workers via shared threshold
         # Determine mode (must match master's logic)
-        enabled_workers = coerce_enabled_worker_ids(enabled_worker_ids)
+        enabled_workers = json.loads(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
@@ -253,47 +181,31 @@ 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 for single-image runs.
-        if num_workers > 0 and batch_size <= 1 and num_tiles_per_image > 1:
+        # and there is more than one tile to process, even if batch == 1.
+        if num_workers > 0 and num_tiles_per_image > 1:
             mode = "static"
             
         debug_log(f"USDU Dist Worker - Batch size {batch_size}")
         
         if mode == "dynamic":
-            return self.process_worker_dynamic(
-                upscaled_image=upscaled_image,
-                core_args=core_args,
-                tile_width=tile_width,
-                tile_height=tile_height,
-                padding=padding,
-                mask_blur=mask_blur,
-                force_uniform_tiles=force_uniform_tiles,
-                multi_job_id=multi_job_id,
-                master_url=master_url,
-                worker_id=worker_id,
-                enabled_worker_ids=enabled_worker_ids,
-                dynamic_threshold=dynamic_threshold,
-            )
+            return self.process_worker_dynamic(upscaled_image, model, positive, negative, vae,
+                                             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_worker_ids, dynamic_threshold)
         
         # Static mode - enhanced with health monitoring and retry logic
-        return self._process_worker_static_sync(upscaled_image, core_args,
+        return self._process_worker_static_sync(upscaled_image, model, positive, negative, vae,
+                                               seed, steps, cfg, sampler_name, scheduler, denoise,
                                                tile_width, tile_height, padding, mask_blur,
-                                               force_uniform_tiles, multi_job_id, master_url,
+                                               force_uniform_tiles, tiled_decode, multi_job_id, master_url,
                                                worker_id, enabled_workers)
 
-    def process_master(
-        self,
-        upscaled_image: Any,
-        core_args: UpscaleCoreArgs,
-        tile_width: int,
-        tile_height: int,
-        padding: int,
-        mask_blur: int,
-        force_uniform_tiles: bool,
-        multi_job_id: str,
-        enabled_worker_ids: str,
-        dynamic_threshold: int,
-    ) -> tuple[Any, ...]:
+    def process_master(self, upscaled_image, model, positive, negative, vae,
+                      seed, steps, cfg, sampler_name, scheduler, denoise,
+                      tile_width, tile_height, padding, mask_blur,
+                      force_uniform_tiles, tiled_decode, multi_job_id, enabled_worker_ids, 
+                      dynamic_threshold):
         """Unified master processing with enhanced monitoring and failure handling."""
         # Round tile dimensions
         tile_width = self.round_to_multiple(tile_width)
@@ -312,206 +224,47 @@ class UltimateSDUpscaleDistributed(
         )
         
         # Parse enabled workers
-        enabled_workers = coerce_enabled_worker_ids(enabled_worker_ids)
+        enabled_workers = json.loads(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,
-        # 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:
+        # 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:
             mode = "static"
         
         log(f"USDU Dist: Workers {num_workers} | Mode {mode} | Threshold {dynamic_threshold}")
 
         if mode == "single_gpu":
             # No workers, process all tiles locally
-            return self.process_single_gpu(
-                upscaled_image=upscaled_image,
-                core_args=core_args,
-                tile_width=tile_width,
-                tile_height=tile_height,
-                padding=padding,
-                mask_blur=mask_blur,
-                force_uniform_tiles=force_uniform_tiles,
-            )
+            return self.process_single_gpu(upscaled_image, model, positive, negative, vae,
+                                         seed, steps, cfg, sampler_name, scheduler, denoise,
+                                         tile_width, tile_height, padding, mask_blur, force_uniform_tiles, tiled_decode)
         
         elif mode == "dynamic":
             # Dynamic mode for large batches
-            return self.process_master_dynamic(
-                upscaled_image=upscaled_image,
-                core_args=core_args,
-                tile_width=tile_width,
-                tile_height=tile_height,
-                padding=padding,
-                mask_blur=mask_blur,
-                force_uniform_tiles=force_uniform_tiles,
-                multi_job_id=multi_job_id,
-                enabled_workers=enabled_workers,
-            )
+            return self.process_master_dynamic(upscaled_image, model, positive, negative, vae,
+                                             seed, steps, cfg, sampler_name, scheduler, denoise,
+                                             tile_width, tile_height, padding, mask_blur,
+                                             force_uniform_tiles, tiled_decode, multi_job_id, enabled_workers)
         
         # Static mode - enhanced with unified job management
-        return self._process_master_static_sync(upscaled_image, core_args,
+        return self._process_master_static_sync(upscaled_image, model, positive, negative, vae,
+                                               seed, steps, cfg, sampler_name, scheduler, denoise,
                                                tile_width, tile_height, padding, mask_blur,
-                                               force_uniform_tiles, multi_job_id, enabled_workers,
+                                               force_uniform_tiles, tiled_decode, multi_job_id, enabled_workers,
                                                all_tiles, num_tiles_per_image)
 
     def _determine_processing_mode(self, batch_size: int, num_workers: int, dynamic_threshold: int) -> str:
-        """Determine mode from worker availability and configured threshold."""
+        """Determines processing mode per requested policy:
+        - any workers     => prefer static (tile-based) for USDU
+        - no workers      => single_gpu
+        """
         if num_workers == 0:
             return "single_gpu"
-        threshold = max(1, int(dynamic_threshold))
-        if int(batch_size) >= threshold:
-            return "dynamic"
+        # Default to static when distributed; master/worker may still override if special cases arise
         return "static"
 
-class USDUDelegateCollector:
-    """Lightweight stand-in used automatically in delegate-only master prompts.
-
-    When the master is in delegate-only mode, the orchestration code swaps the
-    full ``UltimateSDUpscaleDistributed`` class for this one so that no upstream
-    model/image nodes execute on the master.  Workers initialise the dynamic job
-    queue via ``/distributed/init_dynamic_job`` and this node simply waits for
-    all images to arrive, then assembles the output tensor.
-    """
-
-    @classmethod
-    def INPUT_TYPES(cls):
-        return {
-            "required": {},
-            "hidden": {
-                "multi_job_id": ("STRING", {"default": ""}),
-                "enabled_worker_ids": ("STRING", {"default": "[]"}),
-                "delegate_only": ("BOOLEAN", {"default": True}),
-                "is_worker": ("BOOLEAN", {"default": False}),
-                "worker_id": ("STRING", {"default": ""}),
-                "master_url": ("STRING", {"default": ""}),
-            },
-        }
-
-    RETURN_TYPES = ("IMAGE",)
-    FUNCTION = "run"
-    CATEGORY = "image/upscaling"
-
-    @classmethod
-    def IS_CHANGED(cls, **kwargs):
-        return float("nan")
-
-    def run(self, multi_job_id="", enabled_worker_ids="[]", **kwargs):
-        import time
-        import torch
-        from ..utils.image import pil_to_tensor
-        from ..upscale.job_models import ImageJobState
-
-        if not multi_job_id:
-            log("USDUDelegateCollector: no multi_job_id, returning empty image")
-            return (torch.zeros(1, 64, 64, 3),)
-
-        enabled_workers = coerce_enabled_worker_ids(enabled_worker_ids)
-        num_workers = len(enabled_workers)
-        log(f"USDU delegate-only: waiting for workers to init job {multi_job_id}")
-
-        prompt_server = ensure_tile_jobs_initialized()
-
-        # Wait for workers to create the job via /distributed/init_dynamic_job
-        max_wait = 120.0
-        poll_interval = 1.0
-        start = time.time()
-        job_data = None
-        while time.time() - start < max_wait:
-            job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id)
-            if isinstance(job_data, ImageJobState) and job_data.batch_size > 0:
-                break
-            job_data = None
-            time.sleep(poll_interval)
-
-        if job_data is None:
-            log(f"USDUDelegateCollector: job {multi_job_id} was not initialized by workers within {max_wait}s")
-            return (torch.zeros(1, 64, 64, 3),)
-
-        batch_size = job_data.batch_size
-        log(f"USDU delegate-only: job ready, collecting {batch_size} images from {num_workers} workers")
-
-        # Collect all images using the existing result collector
-        collected = run_async_in_server_loop(
-            self._collect_all_images(multi_job_id, batch_size, num_workers, prompt_server),
-            timeout=600.0,
-        )
-
-        # Assemble tensor from collected PIL images
-        result_images = []
-        for idx in range(batch_size):
-            img = collected.get(idx)
-            if img is not None:
-                result_images.append(pil_to_tensor(img))
-            else:
-                log(f"USDUDelegateCollector: missing image {idx}, using blank")
-                result_images.append(torch.zeros(1, 64, 64, 3))
-
-        result_tensor = torch.cat(result_images, dim=0)
-        log(f"USDU delegate-only: collected all {batch_size} images")
-        return (result_tensor,)
-
-    @staticmethod
-    async def _collect_all_images(multi_job_id, batch_size, num_workers, prompt_server):
-        """Wait for all images to arrive from workers."""
-        import asyncio
-        from ..upscale.job_models import ImageJobState
-        from ..upscale.job_timeout import check_and_requeue_timed_out_workers
-
-        timeout_seconds = 300.0
-        poll_interval = 2.0
-        start = asyncio.get_event_loop().time()
-
-        while True:
-            # Drain results from the queue
-            job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id)
-            if isinstance(job_data, ImageJobState):
-                # Drain queue items into completed_images
-                while True:
-                    try:
-                        result = job_data.queue.get_nowait()
-                        worker_id = result.get("worker_id")
-                        if "image_idx" in result and "image" in result:
-                            idx = result["image_idx"]
-                            if idx not in job_data.completed_images:
-                                job_data.completed_images[idx] = result["image"]
-                                debug_log(f"Delegate collected image {idx} from worker {worker_id}")
-                    except asyncio.QueueEmpty:
-                        break
-
-                completed_count = len(job_data.completed_images)
-                if completed_count >= batch_size:
-                    debug_log(f"Delegate: all {batch_size} images collected")
-                    completed = dict(job_data.completed_images)
-                    # Cleanup
-                    async with prompt_server.distributed_tile_jobs_lock:
-                        prompt_server.distributed_pending_tile_jobs.pop(multi_job_id, None)
-                    return completed
-
-                # Check for timeouts periodically
-                await check_and_requeue_timed_out_workers(multi_job_id, batch_size)
-
-            elapsed = asyncio.get_event_loop().time() - start
-            if elapsed > timeout_seconds:
-                log(f"USDUDelegateCollector: timed out after {elapsed:.0f}s with {len(getattr(job_data, 'completed_images', {}))} of {batch_size} images")
-                if isinstance(job_data, ImageJobState):
-                    completed = dict(job_data.completed_images)
-                    async with prompt_server.distributed_tile_jobs_lock:
-                        prompt_server.distributed_pending_tile_jobs.pop(multi_job_id, None)
-                    return completed
-                return {}
-
-            await asyncio.sleep(poll_interval)
-
-
-# Node registration
-NODE_CLASS_MAPPINGS = {
-    "UltimateSDUpscaleDistributed": UltimateSDUpscaleDistributed,
-    "USDUDelegateCollector": USDUDelegateCollector,
-}
-
-NODE_DISPLAY_NAME_MAPPINGS = {
-    "UltimateSDUpscaleDistributed": "Ultimate SD Upscale Distributed (No Upscale)",
-    # No display name for USDUDelegateCollector — it's an internal node
-}
+# Ensure initialization before registering routes
+ensure_tile_jobs_initialized()

nodes/hidden_inputs.py

@@ -1,38 +0,0 @@
-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,
-    assigned_branch_max: int = 9,
-    include_worker_batch_size: bool = False,
-    include_pass_through: bool = False,
-) -> dict[str, tuple[str, dict[str, Any]]]:
-    hidden_inputs: dict[str, tuple[str, dict[str, Any]]] = {
-        "multi_job_id": ("STRING", {"default": ""}),
-        "master_url": ("STRING", {"default": ""}),
-        "delegate_only": ("BOOLEAN", {"default": False}),
-    }
-    hidden_inputs.update(build_worker_identity_hidden_inputs())
-    if include_assigned_branch:
-        hidden_inputs["assigned_branch"] = (
-            "INT",
-            {"default": -1, "min": -1, "max": assigned_branch_max},
-        )
-    if include_worker_batch_size:
-        hidden_inputs["worker_batch_size"] = (
-            "INT",
-            {"default": 1, "min": 1, "max": 1024},
-        )
-    if include_pass_through:
-        hidden_inputs["pass_through"] = ("BOOLEAN", {"default": False})
-    return hidden_inputs

nodes/queue_wait.py

@@ -1,46 +0,0 @@
-import asyncio
-import time
-from collections.abc import Callable
-from typing import Any
-
-from ..utils.config import get_worker_timeout_seconds
-from ..utils.constants import HEARTBEAT_INTERVAL
-from ..utils.logging import log
-
-
-async def collect_worker_queue_results(
-    *,
-    prompt_server: Any,
-    multi_job_id: str,
-    expected_workers: set[str],
-    on_result: Callable[[dict[str, Any]], None],
-    timeout_log_prefix: str,
-    throw_if_interrupted: Callable[[], None],
-) -> set[str]:
-    workers_done: set[str] = set()
-    base_timeout = float(get_worker_timeout_seconds())
-    slice_timeout = min(max(0.1, HEARTBEAT_INTERVAL / 20.0), base_timeout)
-    last_activity = time.time()
-
-    while len(workers_done) < len(expected_workers):
-        throw_if_interrupted()
-        try:
-            async with prompt_server.distributed_jobs_lock:
-                queue = prompt_server.distributed_pending_jobs[multi_job_id]
-            result = await asyncio.wait_for(queue.get(), timeout=slice_timeout)
-        except asyncio.TimeoutError:
-            if (time.time() - last_activity) < base_timeout:
-                continue
-            missing_workers = sorted(expected_workers - workers_done)
-            log(f"{timeout_log_prefix}{missing_workers}")
-            break
-
-        on_result(result)
-        worker_id_value = str(result.get("worker_id", ""))
-        is_last = bool(result.get("is_last", False))
-        last_activity = time.time()
-        base_timeout = float(get_worker_timeout_seconds())
-        if is_last and worker_id_value in expected_workers:
-            workers_done.add(worker_id_value)
-
-    return workers_done

nodes/runtime_helpers.py

@@ -1,16 +0,0 @@
-from typing import Any
-
-
-def get_prompt_server_instance() -> Any:
-    import server as _server
-
-    return _server.PromptServer.instance
-
-
-def throw_if_processing_interrupted() -> None:
-    try:
-        import comfy.model_management as model_management
-
-        model_management.throw_exception_if_processing_interrupted()
-    except Exception:
-        return

nodes/utilities.py

@@ -1,10 +1,7 @@
 import torch
 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]]:
@@ -31,7 +28,7 @@ class DistributedSeed:
     """
     
     @classmethod
-    def INPUT_TYPES(cls: type["DistributedSeed"]) -> dict[str, Any]:
+    def INPUT_TYPES(cls):
         return {
             "required": {
                 "seed": ("INT", {
@@ -42,7 +39,8 @@ class DistributedSeed:
                 }),
             },
             "hidden": {
-                **build_worker_identity_hidden_inputs(),
+                "is_worker": ("BOOLEAN", {"default": False}),
+                "worker_id": ("STRING", {"default": ""}),
             },
         }
     
@@ -51,29 +49,30 @@ class DistributedSeed:
     FUNCTION = "distribute"
     CATEGORY = "utils"
     
-    def distribute(
-        self,
-        seed: int,
-        is_worker: bool = False,
-        worker_id: str = "",
-        enabled_worker_ids: str = "[]",
-    ) -> tuple[int]:
+    def distribute(self, seed, is_worker=False, worker_id=""):
         if not is_worker:
             # Master node: pass through original values
             debug_log(f"Distributor - Master: seed={seed}")
             return (seed,)
         else:
-            enabled_workers = coerce_enabled_worker_ids(enabled_worker_ids)
-            worker_index = parse_worker_index(worker_id, enabled_workers)
-            if worker_index is not None:
+            # 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)
+                
                 offset = worker_index + 1
                 new_seed = seed + offset
                 debug_log(f"Distributor - Worker {worker_index}: seed={seed} → {new_seed}")
                 return (new_seed,)
-
-            debug_log(f"Distributor - Error parsing worker_id '{worker_id}': no worker index resolved")
-            # Fallback: return original seed
-            return (seed,)
+            except (ValueError, IndexError) as e:
+                debug_log(f"Distributor - Error parsing worker_id '{worker_id}': {e}")
+                # Fallback: return original seed
+                return (seed,)
 
 
 # Define ByPassTypeTuple for flexible return types
@@ -93,14 +92,15 @@ class DistributedValue:
     """
 
     @classmethod
-    def INPUT_TYPES(cls: type["DistributedValue"]) -> dict[str, Any]:
+    def INPUT_TYPES(cls):
         return {
             "required": {
                 "default_value": ("STRING", {"default": ""}),
                 "worker_values": ("STRING", {"default": "{}"}),
             },
             "hidden": {
-                **build_worker_identity_hidden_inputs(),
+                "is_worker": ("BOOLEAN", {"default": False}),
+                "worker_id": ("STRING", {"default": ""}),
             },
         }
 
@@ -110,7 +110,7 @@ class DistributedValue:
     CATEGORY = "utils"
 
     @staticmethod
-    def _coerce(value: Any, value_type: str) -> Any:
+    def _coerce(value, value_type):
         """Convert a string value to the requested type."""
         if value_type == "INT":
             return int(float(value))
@@ -119,71 +119,51 @@ class DistributedValue:
         return value  # STRING and COMBO stay as strings
 
     @staticmethod
-    def _coerce_safe(value: Any, value_type: str) -> Any:
+    def _coerce_safe(value, value_type):
         """Best-effort coercion with graceful fallback to original value."""
         try:
             return DistributedValue._coerce(value, value_type)
         except (TypeError, ValueError):
             return value
 
-    @staticmethod
-    def _infer_value_type(value: Any) -> str:
-        """Infer coercion type from the provided default value."""
-        if isinstance(value, bool):
-            return "STRING"
-        if isinstance(value, int):
-            return "INT"
-        if isinstance(value, float):
-            return "FLOAT"
-        return "STRING"
-
-    def distribute(
-        self,
-        default_value: Any,
-        worker_values: str | dict[str, Any] = "{}",
-        is_worker: bool = False,
-        worker_id: str = "",
-        enabled_worker_ids: str = "[]",
-    ) -> tuple[Any]:
+    def distribute(self, default_value, worker_values="{}", is_worker=False, worker_id=""):
         values = {}
         value_type = "STRING"
 
         try:
-            raw_values = json.loads(worker_values) if isinstance(worker_values, str) else worker_values
-            values = dict(raw_values) if isinstance(raw_values, dict) else {}
+            values = json.loads(worker_values) if isinstance(worker_values, str) else worker_values
             if not isinstance(values, dict):
                 values = {}
         except json.JSONDecodeError as e:
             debug_log(f"DistributedValue - Error parsing worker_values: {e}")
             values = {}
 
-        inferred_type = self._infer_value_type(default_value)
-        value_type = values.get("_type", inferred_type)
-        if value_type not in {"STRING", "COMBO", "INT", "FLOAT"}:
-            value_type = inferred_type
+        value_type = values.get("_type", "STRING")
         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,)
 
-        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:
+        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:
                 coerced = self._coerce(raw, value_type)
-                debug_log(f"DistributedValue - Worker key {lookup_key}: returning '{coerced}'")
+                debug_log(f"DistributedValue - Worker {idx}: returning '{coerced}'")
                 return (coerced,)
-            except (TypeError, ValueError) as e:
-                debug_log(f"DistributedValue - Error coercing worker value for key {lookup_key}: {e}")
+        except (ValueError, IndexError) as e:
+            debug_log(f"DistributedValue - Error: {e}")
         debug_log(f"DistributedValue - Worker fallback: returning default '{coerced_default}'")
         return (coerced_default,)
 
 class DistributedModelName:
     @classmethod
-    def INPUT_TYPES(cls: type["DistributedModelName"]) -> dict[str, Any]:
+    def INPUT_TYPES(cls):
         return {
             "required": {
                 "text": ("STRING", {"default": ""}),
@@ -200,7 +180,7 @@ class DistributedModelName:
     OUTPUT_NODE = True
     CATEGORY = "utils"
 
-    def _stringify(self, value: Any) -> str:
+    def _stringify(self, value):
         if isinstance(value, str):
             return value
         if isinstance(value, (int, float, bool)):
@@ -210,7 +190,7 @@ class DistributedModelName:
         except Exception:
             return str(value)
 
-    def _update_workflow(self, extra_pnginfo: Any, unique_id: Any, values: list[str]) -> None:
+    def _update_workflow(self, extra_pnginfo, unique_id, values):
         if not extra_pnginfo:
             return
         info = extra_pnginfo[0] if isinstance(extra_pnginfo, list) else extra_pnginfo
@@ -228,12 +208,7 @@ class DistributedModelName:
         if node:
             node["widgets_values"] = [values]
 
-    def log_input(
-        self,
-        text: Any,
-        unique_id: Any = None,
-        extra_pnginfo: Any = None,
-    ) -> dict[str, Any]:
+    def log_input(self, text, unique_id=None, extra_pnginfo=None):
         values = []
         if isinstance(text, list):
             for val in text:
@@ -249,7 +224,7 @@ class DistributedModelName:
         return {"ui": {"text": values}, "result": (values,)}
 
 class ByPassTypeTuple(tuple):
-    def __getitem__(self, index: int) -> Any:
+    def __getitem__(self, index):
         if index > 0:
             index = 0
         item = super().__getitem__(index)
@@ -259,7 +234,7 @@ class ByPassTypeTuple(tuple):
 
 class ImageBatchDivider:
     @classmethod
-    def INPUT_TYPES(cls: type["ImageBatchDivider"]) -> dict[str, Any]:
+    def INPUT_TYPES(s):
         return {
             "required": {
                 "images": ("IMAGE",),
@@ -280,7 +255,7 @@ class ImageBatchDivider:
     OUTPUT_NODE = True
     CATEGORY = "image"
     
-    def divide_batch(self, images: torch.Tensor, divide_by: int) -> tuple[torch.Tensor, ...]:
+    def divide_batch(self, images, divide_by):
         total_splits = max(1, min(int(divide_by), 10))
         total_frames = images.shape[0]
         empty_tensor = images[:0]
@@ -294,10 +269,10 @@ class ImageBatchDivider:
 
 
 class AudioBatchDivider:
-    """Divides an audio waveform into multiple parts along the time/samples dimension."""
+    """Divides an audio waveform into sequential segments along the time dimension."""
 
     @classmethod
-    def INPUT_TYPES(cls: type["AudioBatchDivider"]) -> dict[str, Any]:
+    def INPUT_TYPES(s):
         return {
             "required": {
                 "audio": ("AUDIO",),
@@ -307,7 +282,7 @@ class AudioBatchDivider:
                     "max": 10,
                     "step": 1,
                     "display": "number",
-                    "tooltip": "Number of parts to divide the audio into"
+                    "tooltip": "Number of sequential time segments to create"
                 }),
             }
         }
@@ -318,7 +293,7 @@ class AudioBatchDivider:
     OUTPUT_NODE = True
     CATEGORY = "audio"
 
-    def divide_audio(self, audio: dict[str, Any], divide_by: int) -> tuple[dict[str, Any], ...]:
+    def divide_audio(self, audio, divide_by):
         import torch
 
         waveform = audio.get("waveform")
@@ -358,7 +333,7 @@ class DistributedEmptyImage:
     """Produces an empty IMAGE batch used when the master delegates all work."""
 
     @classmethod
-    def INPUT_TYPES(cls: type["DistributedEmptyImage"]) -> dict[str, Any]:
+    def INPUT_TYPES(cls):
         return {
             "required": {
                 "height": ("INT", {"default": 64, "min": 1, "max": 4096, "step": 1}),
@@ -371,7 +346,7 @@ class DistributedEmptyImage:
     FUNCTION = "create"
     CATEGORY = "image"
 
-    def create(self, height: int, width: int, channels: int) -> tuple[torch.Tensor]:
+    def create(self, height, width, channels):
         import torch
 
         shape = (0, height, width, channels)

nodes/v3.py

@@ -0,0 +1,214 @@
+"""Native V3 schemas with the existing execution algorithms kept intact.
+
+Runtime objects are private, per-execution helpers, not registered V1 nodes.
+This avoids sharing mutable instance state through V3's sanitized class clones.
+"""
+import comfy.samplers
+from comfy_api.v0_0_2 import io
+
+from . import utilities as _utilities
+from .collector import DistributedCollectorNode as _CollectorRuntime
+from ..utils import transfer_codec as _transfer_codec
+from .distributed_upscale import UltimateSDUpscaleDistributed as _UpscaleRuntime
+
+
+class DistributedSeed(io.ComfyNode):
+    @classmethod
+    def define_schema(cls):
+        return io.Schema(
+            node_id='DistributedSeed', display_name='Distributed Seed', category='utils',
+            inputs=[io.Int.Input('seed', default=1125899906842, min=0,
+                                 max=1125899906842624, force_input=False)],
+            outputs=[io.Int.Output(display_name='seed')],
+            accept_all_inputs=True,
+        )
+
+    @classmethod
+    def execute(cls, seed, is_worker=False, worker_id=''):
+        return io.NodeOutput(*_utilities.DistributedSeed().distribute(seed, is_worker, worker_id))
+
+
+class DistributedValue(io.ComfyNode):
+    @classmethod
+    def define_schema(cls):
+        return io.Schema(
+            node_id='DistributedValue', display_name='Distributed Value', category='utils',
+            inputs=[io.String.Input('default_value', default=''),
+                    io.String.Input('worker_values', default='{}')],
+            outputs=[io.AnyType.Output(display_name='value')],
+            accept_all_inputs=True,
+        )
+
+    @classmethod
+    def execute(cls, default_value, worker_values='{}', is_worker=False, worker_id=''):
+        return io.NodeOutput(*_utilities.DistributedValue().distribute(
+            default_value, worker_values, is_worker, worker_id))
+
+
+class DistributedModelName(io.ComfyNode):
+    @classmethod
+    def define_schema(cls):
+        return io.Schema(
+            node_id='DistributedModelName', display_name='Distributed Model Name', category='utils',
+            inputs=[io.String.Input('text', default='')],
+            outputs=[io.AnyType.Output(display_name='output')],
+            hidden=[io.Hidden.unique_id, io.Hidden.extra_pnginfo], is_output_node=True,
+        )
+
+    @classmethod
+    def execute(cls, text):
+        result = _utilities.DistributedModelName().log_input(
+            text, unique_id=cls.hidden.unique_id, extra_pnginfo=cls.hidden.extra_pnginfo)
+        return io.NodeOutput(*result['result'], ui=result['ui'])
+
+
+class ImageBatchDivider(io.ComfyNode):
+    @classmethod
+    def define_schema(cls):
+        return io.Schema(
+            node_id='ImageBatchDivider', display_name='Image Batch Divider', category='image',
+            inputs=[io.Image.Input('images'),
+                    io.Int.Input('divide_by', default=2, min=1, max=10, step=1,
+                                 display_mode=io.NumberDisplay.number,
+                                 tooltip='Number of parts to divide the batch into')],
+            # The existing frontend still displays only divide_by sockets.
+            outputs=[io.Image.Output(display_name=f'batch_{index + 1}') for index in range(10)],
+            is_output_node=True,
+        )
+
+    @classmethod
+    def execute(cls, images, divide_by):
+        return io.NodeOutput(*_utilities.ImageBatchDivider().divide_batch(images, divide_by))
+
+
+class AudioBatchDivider(io.ComfyNode):
+    @classmethod
+    def define_schema(cls):
+        return io.Schema(
+            node_id='AudioBatchDivider', display_name='Audio Segment Divider', category='audio',
+            inputs=[io.Audio.Input('audio'),
+                    io.Int.Input('divide_by', default=2, min=1, max=10, step=1,
+                                 display_mode=io.NumberDisplay.number,
+                                 tooltip='Number of sequential time segments to create')],
+            outputs=[io.Audio.Output(display_name=f'audio_{index + 1}') for index in range(10)],
+            is_output_node=True,
+        )
+
+    @classmethod
+    def execute(cls, audio, divide_by):
+        return io.NodeOutput(*_utilities.AudioBatchDivider().divide_audio(audio, divide_by))
+
+
+class DistributedEmptyImage(io.ComfyNode):
+    @classmethod
+    def define_schema(cls):
+        return io.Schema(
+            node_id='DistributedEmptyImage', display_name='Distributed Empty Image', category='image',
+            inputs=[io.Int.Input('height', default=64, min=1, max=4096, step=1),
+                    io.Int.Input('width', default=64, min=1, max=4096, step=1),
+                    io.Int.Input('channels', default=3, min=1, max=4, step=1)],
+            outputs=[io.Image.Output()],
+        )
+
+    @classmethod
+    def execute(cls, height, width, channels):
+        return io.NodeOutput(*_utilities.DistributedEmptyImage().create(height, width, channels))
+
+
+class DistributedCollector(io.ComfyNode):
+    @classmethod
+    def define_schema(cls):
+        return io.Schema(
+            node_id='DistributedCollector', display_name='Distributed Collector', category='image',
+            inputs=[io.Boolean.Input('load_balance', default=False,
+                                     tooltip='Run this workflow on one least-busy participant (master included when participating).'),
+                    io.Image.Input('images', optional=True), io.Audio.Input('audio', optional=True),
+                    io.Combo.Input('transfer_format', options=list(_transfer_codec.TRANSFER_FORMATS),
+                                   default=_transfer_codec.DEFAULT_TRANSFER_FORMAT, optional=True,
+                                   tooltip=_transfer_codec.TRANSFER_FORMAT_TOOLTIP),
+                    io.Int.Input('transfer_quality', default=_transfer_codec.DEFAULT_TRANSFER_QUALITY,
+                                 min=1, max=100, step=1, optional=True,
+                                 tooltip=_transfer_codec.TRANSFER_QUALITY_TOOLTIP)],
+            outputs=[io.Image.Output(display_name='images'), io.Audio.Output(display_name='audio')],
+            is_input_list=True, accept_all_inputs=True,
+        )
+
+    @classmethod
+    def execute(cls, images=None, load_balance=False, audio=None, multi_job_id='',
+                is_worker=False, master_url='', enabled_worker_ids='[]', worker_batch_size=1,
+                worker_id='', pass_through=False, delegate_only=False,
+                transfer_format=_transfer_codec.DEFAULT_TRANSFER_FORMAT,
+                transfer_quality=_transfer_codec.DEFAULT_TRANSFER_QUALITY):
+        return io.NodeOutput(*_CollectorRuntime().run(
+            images=images, load_balance=load_balance, audio=audio, multi_job_id=multi_job_id,
+            is_worker=is_worker, master_url=master_url, enabled_worker_ids=enabled_worker_ids,
+            worker_batch_size=worker_batch_size, worker_id=worker_id,
+            pass_through=pass_through, delegate_only=delegate_only,
+            transfer_format=transfer_format, transfer_quality=transfer_quality))
+
+
+class UltimateSDUpscaleDistributed(io.ComfyNode):
+    @classmethod
+    def define_schema(cls):
+        return io.Schema(
+            node_id='UltimateSDUpscaleDistributed',
+            display_name='Ultimate SD Upscale Distributed (No Upscale)', category='image/upscaling',
+            inputs=[
+                io.Image.Input('upscaled_image'), io.Model.Input('model'),
+                io.Conditioning.Input('positive'), io.Conditioning.Input('negative'), io.Vae.Input('vae'),
+                io.Int.Input('seed', default=0, min=0, max=0xffffffffffffffff),
+                io.Int.Input('steps', default=20, min=1, max=10000),
+                io.Float.Input('cfg', default=8.0, min=0.0, max=100.0),
+                io.Combo.Input('sampler_name', options=comfy.samplers.KSampler.SAMPLERS),
+                io.Combo.Input('scheduler', options=comfy.samplers.KSampler.SCHEDULERS),
+                io.Float.Input('denoise', default=0.5, min=0.0, max=1.0, step=0.01),
+                io.Int.Input('tile_width', default=512, min=64, max=2048, step=8),
+                io.Int.Input('tile_height', default=512, min=64, max=2048, step=8),
+                io.Int.Input('padding', default=32, min=0, max=256, step=8),
+                io.Int.Input('mask_blur', default=8, min=0, max=256),
+                io.Boolean.Input('force_uniform_tiles', default=True),
+                io.Boolean.Input('tiled_decode', default=False),
+            ], outputs=[io.Image.Output()], accept_all_inputs=True,
+        )
+
+    @classmethod
+    def fingerprint_inputs(cls, **kwargs):
+        return _UpscaleRuntime.IS_CHANGED(**kwargs)
+
+    @classmethod
+    def execute(cls, upscaled_image, model, positive, negative, vae, seed, steps, cfg,
+                sampler_name, scheduler, denoise, tile_width, tile_height, padding,
+                mask_blur, force_uniform_tiles, tiled_decode, multi_job_id='', is_worker=False,
+                master_url='', enabled_worker_ids='[]', worker_id='', tile_indices='', dynamic_threshold=8):
+        return io.NodeOutput(*_UpscaleRuntime().run(
+            upscaled_image, model, positive, negative, vae, seed, steps, cfg,
+            sampler_name, scheduler, denoise, tile_width, tile_height, padding,
+            mask_blur, force_uniform_tiles, tiled_decode, multi_job_id, is_worker,
+            master_url, enabled_worker_ids, worker_id, tile_indices, dynamic_threshold))
+
+
+class DistributedWorkerSink(io.ComfyNode):
+    """Internal output node that terminates pruned worker prompts.
+
+    Worker prompts drop everything downstream of the distributed node, which
+    leaves no output node for ComfyUI to execute towards. This sink provides
+    one without writing preview files (PreviewImage wrote every frame to the
+    worker's temp directory). Hidden from search/menus; inserted by the master.
+    """
+
+    @classmethod
+    def define_schema(cls):
+        return io.Schema(
+            node_id='DistributedWorkerSink', display_name='Distributed Worker Sink (internal)', category='utils',
+            inputs=[io.AnyType.Input('value')],
+            is_output_node=True, is_dev_only=True,
+        )
+
+    @classmethod
+    def execute(cls, value):
+        return io.NodeOutput()
+
+
+NODES = [DistributedCollector, DistributedSeed, DistributedModelName, DistributedValue,
+         ImageBatchDivider, AudioBatchDivider, DistributedEmptyImage, UltimateSDUpscaleDistributed,
+         DistributedWorkerSink]

pyproject.toml

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

runtime/__init__.py

@@ -0,0 +1 @@
+"""Internal extension lifecycle support (not a node provider)."""

runtime/bootstrap.py

@@ -0,0 +1,35 @@
+"""Initialize routes, distributed state and the existing worker lifecycle."""
+import atexit
+import os
+
+import server
+
+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
+from ..upscale.job_store import ensure_tile_jobs_initialized
+
+_initialized = False
+
+
+def initialize():
+    """Called by ComfyExtension.on_load; initialize once per loaded package."""
+    global _initialized
+    if _initialized:
+        return
+
+    ensure_config_exists()
+    from .. import api  # noqa: F401 - registers the existing @routes.* handlers
+    from ..api.queue_orchestration import ensure_distributed_state
+
+    ensure_distributed_state(server.PromptServer.instance)
+    ensure_tile_jobs_initialized()
+
+    if not os.environ.get('COMFYUI_IS_WORKER'):
+        atexit.register(sync_cleanup)
+        delayed_auto_launch()
+        register_async_signals()
+
+    _initialized = True
+    debug_log('Loaded Distributed nodes.')
+    debug_log(f'Config file: {CONFIG_FILE}')

tests/api/harness.py

@@ -1,163 +0,0 @@
-"""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

@@ -1,423 +0,0 @@
-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,19 +3,9 @@ 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):
@@ -35,18 +25,38 @@ def _load_config_routes_module():
     module_path = Path(__file__).resolve().parents[2] / "api" / "config_routes.py"
     package_name = "dist_api_config_testpkg"
 
-    bootstrap_test_package(package_name, with_api=True, with_utils=True)
+    for mod_name in list(sys.modules):
+        if mod_name == package_name or mod_name.startswith(f"{package_name}."):
+            del sys.modules[mod_name]
 
-    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
+    root_pkg = types.ModuleType(package_name)
+    root_pkg.__path__ = []
+    sys.modules[package_name] = root_pkg
 
-    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()
+    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__ = [str(module_path.parents[1] / "utils")]
+    sys.modules[f"{package_name}.utils"] = utils_pkg
+
+    import aiohttp  # Load the real dependency; isolate the fake response at the tested module.
+
+    class _Routes:
+        def get(self, _path):
+            def _decorator(fn):
+                return fn
+            return _decorator
+
+        def post(self, _path):
+            def _decorator(fn):
+                return fn
+            return _decorator
+
+    server_module = types.ModuleType("server")
+    server_module.PromptServer = types.SimpleNamespace(instance=types.SimpleNamespace(routes=_Routes()))
+    sys.modules["server"] = server_module
 
     logging_module = types.ModuleType(f"{package_name}.utils.logging")
     logging_module.debug_log = lambda *_args, **_kwargs: None
@@ -56,16 +66,7 @@ def _load_config_routes_module():
     network_module = types.ModuleType(f"{package_name}.utils.network")
 
     async def _handle_api_error(_request, error, status=500):
-        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,
-        )
+        return _FakeResponse({"status": "error", "message": str(error)}, status=status)
 
     network_module.handle_api_error = _handle_api_error
     network_module.normalize_host = lambda value: value
@@ -81,12 +82,6 @@ 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)
@@ -94,7 +89,14 @@ def _load_config_routes_module():
     assert spec is not None and spec.loader is not None
     spec.loader.exec_module(module)
 
-    cleanup_optional_module("aiohttp", created_aiohttp_stub)
+    class _EmptySecretStore:
+        def has(self, _worker_id):
+            return False
+
+    module.web = types.SimpleNamespace(
+        json_response=lambda payload, status=200: _FakeResponse(payload, status=status)
+    )
+    module._serverless_secret_store = lambda: _EmptySecretStore()
 
     return module
 
@@ -115,7 +117,9 @@ 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):
+        with patch.object(config_routes, "load_config", return_value=cfg), patch.object(
+            config_routes, "save_config", return_value=True
+        ):
             response = await config_routes.update_config_endpoint(_FakeRequest({"debug": True}))
 
         self.assertEqual(response.status, 200)

tests/api/test_distributed_queue.py

@@ -0,0 +1,555 @@
+import importlib.util
+import sys
+import types
+import unittest
+import asyncio
+import base64
+import io
+import json
+from dataclasses import dataclass
+from pathlib import Path
+from unittest.mock import AsyncMock, patch
+
+import numpy as np
+import torch
+
+
+class _FakeResponse:
+    def __init__(self, payload, status=200):
+        self.payload = payload
+        self.status = status
+
+
+class _FakeRequest:
+    def __init__(self, payload):
+        self._payload = payload
+
+    async def json(self):
+        return self._payload
+
+
+def _register_real_transfer_codec(package_name, utils_pkg):
+    """Load the real utils/transfer_codec.py into the synthetic test package."""
+    codec_path = Path(__file__).resolve().parents[2] / "utils" / "transfer_codec.py"
+    codec_spec = importlib.util.spec_from_file_location(f"{package_name}.utils.transfer_codec", codec_path)
+    codec_module = importlib.util.module_from_spec(codec_spec)
+    sys.modules[f"{package_name}.utils.transfer_codec"] = codec_module
+    codec_spec.loader.exec_module(codec_module)
+    utils_pkg.transfer_codec = codec_module
+    return codec_module
+
+
+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
+    _register_real_transfer_codec(package_name, utils_pkg)
+
+    # aiohttp.web stub
+    created_aiohttp_stub = False
+    if "aiohttp" not in sys.modules:
+        created_aiohttp_stub = True
+        aiohttp_module = types.ModuleType("aiohttp")
+        aiohttp_module.web = types.SimpleNamespace(
+            json_response=lambda payload, status=200: _FakeResponse(payload, status=status)
+        )
+        sys.modules["aiohttp"] = aiohttp_module
+
+    # server module stub with route decorators
+    class _Routes:
+        def get(self, _path):
+            def _decorator(fn):
+                return fn
+            return _decorator
+
+        def post(self, _path):
+            def _decorator(fn):
+                return fn
+            return _decorator
+
+    prompt_server_instance = types.SimpleNamespace(
+        routes=_Routes(),
+        distributed_jobs_lock=None,
+        distributed_pending_jobs={},
+    )
+    server_module = types.ModuleType("server")
+    server_module.PromptServer = types.SimpleNamespace(instance=prompt_server_instance)
+    sys.modules["server"] = server_module
+
+    # torch stub (only needed to satisfy import)
+    created_torch_stub = False
+    if "torch" not in sys.modules:
+        created_torch_stub = True
+        torch_module = types.ModuleType("torch")
+        torch_module.cuda = types.SimpleNamespace(
+            is_available=lambda: False,
+            empty_cache=lambda: None,
+            ipc_collect=lambda: None,
+        )
+        sys.modules["torch"] = torch_module
+
+    # PIL stub (only needed to satisfy import)
+    created_pil_stub = False
+    if "PIL" not in sys.modules:
+        created_pil_stub = True
+        pil_module = types.ModuleType("PIL")
+        image_module = types.ModuleType("PIL.Image")
+        pil_module.Image = image_module
+        sys.modules["PIL"] = pil_module
+        sys.modules["PIL.Image"] = image_module
+
+    logging_module = types.ModuleType(f"{package_name}.utils.logging")
+    logging_module.debug_log = lambda *_args, **_kwargs: None
+    logging_module.log = lambda *_args, **_kwargs: None
+    sys.modules[f"{package_name}.utils.logging"] = logging_module
+
+    image_module = types.ModuleType(f"{package_name}.utils.image")
+    image_module.pil_to_tensor = lambda *_args, **_kwargs: None
+    image_module.ensure_contiguous = lambda tensor: tensor
+    sys.modules[f"{package_name}.utils.image"] = image_module
+
+    audio_payload_module = types.ModuleType(f"{package_name}.utils.audio_payload")
+
+    def _decode_audio_payload(payload):
+        if payload is None:
+            return None
+        if not isinstance(payload, dict):
+            raise ValueError("Field 'audio' must be an object when provided.")
+
+        encoded = payload.get("data")
+        shape = payload.get("shape")
+        dtype = payload.get("dtype", "float32")
+        sample_rate = payload.get("sample_rate", 44100)
+        if not isinstance(encoded, str) or not encoded.strip():
+            raise ValueError("Field 'audio.data' must be a non-empty base64 string.")
+        if not isinstance(shape, list) or len(shape) != 3:
+            raise ValueError("Field 'audio.shape' must be a 3-item list.")
+        if dtype != "float32":
+            raise ValueError("Field 'audio.dtype' must be 'float32'.")
+        try:
+            shape_tuple = tuple(int(dim) for dim in shape)
+        except Exception as exc:
+            raise ValueError("Field 'audio.shape' must contain integers.") from exc
+
+        raw = base64.b64decode(encoded, validate=True)
+        expected_bytes = int(np.prod(shape_tuple, dtype=np.int64)) * 4
+        if len(raw) != expected_bytes:
+            raise ValueError("Field 'audio.data' byte size mismatch.")
+
+        waveform = torch.from_numpy(np.frombuffer(raw, dtype=np.float32).reshape(shape_tuple).copy())
+        return {"waveform": waveform, "sample_rate": int(sample_rate)}
+
+    audio_payload_module.decode_audio_payload = _decode_audio_payload
+    sys.modules[f"{package_name}.utils.audio_payload"] = audio_payload_module
+
+    network_module = types.ModuleType(f"{package_name}.utils.network")
+
+    async def _handle_api_error(_request, error, status=500):
+        return _FakeResponse({"status": "error", "message": str(error)}, status=status)
+
+    network_module.handle_api_error = _handle_api_error
+    sys.modules[f"{package_name}.utils.network"] = network_module
+
+    constants_module = types.ModuleType(f"{package_name}.utils.constants")
+    constants_module.MEMORY_CLEAR_DELAY = 0.0
+    constants_module.JOB_INIT_GRACE_PERIOD = 10.0
+    sys.modules[f"{package_name}.utils.constants"] = constants_module
+
+    async_helpers_module = types.ModuleType(f"{package_name}.utils.async_helpers")
+    async_helpers_module.queue_prompt_payload = AsyncMock(return_value="prompt_local")
+    sys.modules[f"{package_name}.utils.async_helpers"] = async_helpers_module
+
+    queue_orchestration_module = types.ModuleType(f"{package_name}.api.queue_orchestration")
+    queue_orchestration_module.orchestrate_distributed_execution = AsyncMock(return_value=("prompt_dist", 7, 1, {}))
+    sys.modules[f"{package_name}.api.queue_orchestration"] = queue_orchestration_module
+
+    @dataclass(frozen=True)
+    class _QueuePayload:
+        prompt: dict
+        workflow_meta: object
+        client_id: str
+        delegate_master: object
+        enabled_worker_ids: list
+        auto_prepare: bool
+        trace_execution_id: object
+
+    def _parse_queue_request_payload(data):
+        if not isinstance(data, dict):
+            raise ValueError("Expected a JSON object body")
+        prompt = data.get("prompt")
+        if not isinstance(prompt, dict):
+            raise ValueError("Field 'prompt' must be an object")
+        enabled = data.get("enabled_worker_ids")
+        if not isinstance(enabled, list):
+            raise ValueError("enabled_worker_ids required")
+        client_id = data.get("client_id")
+        if not isinstance(client_id, str) or not client_id.strip():
+            raise ValueError("client_id required")
+        return _QueuePayload(
+            prompt=prompt,
+            workflow_meta=data.get("workflow"),
+            client_id=client_id,
+            delegate_master=data.get("delegate_master"),
+            enabled_worker_ids=enabled,
+            auto_prepare=bool(data.get("auto_prepare", True)),
+            trace_execution_id=data.get("trace_execution_id"),
+        )
+
+    queue_request_module = types.ModuleType(f"{package_name}.api.queue_request")
+    queue_request_module.parse_queue_request_payload = _parse_queue_request_payload
+    sys.modules[f"{package_name}.api.queue_request"] = queue_request_module
+
+    spec = importlib.util.spec_from_file_location(f"{package_name}.api.job_routes", module_path)
+    module = importlib.util.module_from_spec(spec)
+    assert spec is not None and spec.loader is not None
+    spec.loader.exec_module(module)
+    # Keep response assertions independent of real aiohttp import ordering.
+    module.web = types.SimpleNamespace(
+        json_response=lambda payload, status=200: _FakeResponse(payload, status=status)
+    )
+
+    if created_aiohttp_stub:
+        sys.modules.pop("aiohttp", None)
+    if created_torch_stub:
+        sys.modules.pop("torch", None)
+    if created_pil_stub:
+        sys.modules.pop("PIL.Image", None)
+        sys.modules.pop("PIL", None)
+
+    return module
+
+
+job_routes = _load_job_routes_module()
+
+
+class DistributedQueueEndpointTests(unittest.IsolatedAsyncioTestCase):
+    async def test_distributed_queue_happy_path_returns_prompt_metadata(self):
+        request = _FakeRequest(
+            {
+                "prompt": {"1": {"class_type": "Node"}},
+                "enabled_worker_ids": ["w1"],
+                "client_id": "client-1",
+                "auto_prepare": True,
+            }
+        )
+        with patch.object(
+            job_routes,
+            "orchestrate_distributed_execution",
+            new=AsyncMock(return_value=("prompt_123", 42, 2, {})),
+        ):
+            response = await job_routes.distributed_queue_endpoint(request)
+
+        self.assertEqual(response.status, 200)
+        self.assertEqual(response.payload.get("prompt_id"), "prompt_123")
+        self.assertEqual(response.payload.get("number"), 42)
+        self.assertEqual(response.payload.get("node_errors"), {})
+        self.assertTrue(response.payload.get("auto_prepare_supported"))
+
+    async def test_distributed_queue_missing_prompt_returns_400(self):
+        request = _FakeRequest(
+            {
+                "enabled_worker_ids": ["w1"],
+                "client_id": "client-1",
+            }
+        )
+        response = await job_routes.distributed_queue_endpoint(request)
+        self.assertEqual(response.status, 400)
+        self.assertIn("prompt", response.payload.get("message", "").lower())
+
+    async def test_distributed_queue_missing_enabled_worker_ids_returns_400(self):
+        request = _FakeRequest(
+            {
+                "prompt": {"1": {"class_type": "Node"}},
+                "client_id": "client-1",
+            }
+        )
+        response = await job_routes.distributed_queue_endpoint(request)
+        self.assertEqual(response.status, 400)
+        self.assertIn("enabled_worker_ids", response.payload.get("message", "").lower())
+
+
+class JobCompleteAudioPayloadTests(unittest.IsolatedAsyncioTestCase):
+    def _encoded_audio_payload(self):
+        waveform = np.arange(8, dtype=np.float32).reshape(1, 2, 4)
+        return {
+            "sample_rate": 44100,
+            "shape": [1, 2, 4],
+            "dtype": "float32",
+            "data": base64.b64encode(waveform.tobytes()).decode("ascii"),
+        }
+
+    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}
+        request = _FakeRequest(
+            {
+                "job_id": "job-1",
+                "worker_id": "worker-1",
+                "batch_idx": 0,
+                "image": "data:image/png;base64,AAAA",
+                "audio": self._encoded_audio_payload(),
+                "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, 200)
+        queued = await queue.get()
+        self.assertEqual(queued["worker_id"], "worker-1")
+        self.assertTrue(queued["is_last"])
+        self.assertIsNotNone(queued["audio"])
+        self.assertEqual(queued["audio"]["sample_rate"], 44100)
+        self.assertEqual(tuple(queued["audio"]["waveform"].shape), (1, 2, 4))
+
+    async def test_job_complete_accepts_audio_without_image(self):
+        queue = asyncio.Queue()
+        job_routes.prompt_server.distributed_jobs_lock = asyncio.Lock()
+        job_routes.prompt_server.distributed_pending_jobs = {"audio-only-job": queue}
+        request = _FakeRequest(
+            {
+                "job_id": "audio-only-job",
+                "worker_id": "worker-1",
+                "batch_idx": 0,
+                "audio": self._encoded_audio_payload(),
+                "is_last": True,
+            }
+        )
+
+        with patch.object(job_routes, "_decode_canonical_png_tensor") as decode_image:
+            response = await job_routes.job_complete_endpoint(request)
+
+        self.assertEqual(response.status, 200)
+        decode_image.assert_not_called()
+        queued = await queue.get()
+        self.assertIsNone(queued["tensor"])
+        self.assertEqual(queued["audio"]["sample_rate"], 44100)
+
+    async def test_job_complete_rejects_payload_without_image_or_audio(self):
+        request = _FakeRequest(
+            {
+                "job_id": "job-1",
+                "worker_id": "worker-1",
+                "batch_idx": 0,
+                "is_last": True,
+            }
+        )
+
+        response = await job_routes.job_complete_endpoint(request)
+
+        self.assertEqual(response.status, 400)
+        self.assertIn("image or audio", response.payload.get("message", "").lower())
+
+    async def test_job_complete_rejects_invalid_image_even_when_audio_is_present(self):
+        request = _FakeRequest(
+            {
+                "job_id": "job-1",
+                "worker_id": "worker-1",
+                "batch_idx": 0,
+                "image": 123,
+                "audio": self._encoded_audio_payload(),
+                "is_last": True,
+            }
+        )
+
+        response = await job_routes.job_complete_endpoint(request)
+
+        self.assertEqual(response.status, 400)
+        self.assertIn("image", response.payload.get("message", "").lower())
+
+    async def test_job_complete_returns_400_for_malformed_audio(self):
+        request = _FakeRequest(
+            {
+                "job_id": "job-1",
+                "worker_id": "worker-1",
+                "batch_idx": 0,
+                "audio": {"data": "AAAA", "shape": [1, 2], "dtype": "float32"},
+                "is_last": True,
+            }
+        )
+
+        response = await job_routes.job_complete_endpoint(request)
+
+        self.assertEqual(response.status, 400)
+        self.assertIn("audio.shape", response.payload.get("message", "").lower())
+
+    def test_decode_audio_payload_rejects_bad_shape(self):
+        bad = {
+            "sample_rate": 44100,
+            "shape": [1, 2],
+            "dtype": "float32",
+            "data": base64.b64encode(b"\x00\x00\x00\x00").decode("ascii"),
+        }
+        with self.assertRaises(ValueError):
+            job_routes._decode_audio_payload(bad)
+
+    def test_decode_audio_payload_rejects_bad_dtype(self):
+        payload = {
+            "sample_rate": 44100,
+            "shape": [1, 2, 4],
+            "dtype": "float16",
+            "data": base64.b64encode((np.zeros((1, 2, 4), dtype=np.float32)).tobytes()).decode("ascii"),
+        }
+        with self.assertRaises(ValueError):
+            job_routes._decode_audio_payload(payload)
+
+
+class _FakeFileField:
+    def __init__(self, data):
+        self.file = io.BytesIO(data)
+
+
+class _FakeMultipartRequest:
+    content_type = "multipart/form-data"
+
+    def __init__(self, envelope, parts):
+        from multidict import MultiDict
+
+        fields = [("envelope", envelope if isinstance(envelope, str) else json.dumps(envelope))]
+        fields += [("media", _FakeFileField(part)) for part in parts]
+        self._form = MultiDict(fields)
+
+    async def post(self):
+        return self._form
+
+    async def json(self):
+        raise ValueError("not JSON")
+
+
+class JobCompleteCompressedPayloadTests(unittest.IsolatedAsyncioTestCase):
+    codec = job_routes.transfer_codec
+
+    def setUp(self):
+        self.queue = asyncio.Queue()
+        job_routes.prompt_server.distributed_jobs_lock = asyncio.Lock()
+        job_routes.prompt_server.distributed_pending_jobs = {"job-z": self.queue}
+
+    def _frames(self, count=4, height=20, width=30):
+        rng = np.random.default_rng(1)
+        return rng.integers(0, 256, (count, height, width, 3), dtype=np.uint8)
+
+    def _request_for(self, frames, fmt, *, start_offset=0, is_last=True, audio=None, **overrides):
+        chunk = next(self.codec.encode_chunks(frames, fmt, 90, 10**9))
+        parts = chunk.pop("parts")
+        envelope = {
+            "job_id": "job-z",
+            "worker_id": "worker-1",
+            "is_last": is_last,
+            **chunk,
+            "start_index": chunk["start_index"] + start_offset,
+        }
+        if audio is not None:
+            envelope["audio"] = audio
+        envelope.update(overrides)
+        return _FakeMultipartRequest(envelope, parts)
+
+    async def _drain(self):
+        items = []
+        while not self.queue.empty():
+            items.append(await self.queue.get())
+        return items
+
+    async def test_lossless_chunk_enqueues_one_item_per_frame_bit_exact(self):
+        frames = self._frames()
+        response = await job_routes.job_complete_endpoint(self._request_for(frames, "ffv1", start_offset=10))
+
+        self.assertEqual(response.status, 200)
+        items = await self._drain()
+        self.assertEqual([item["image_index"] for item in items], [10, 11, 12, 13])
+        self.assertEqual([item["is_last"] for item in items], [False, False, False, True])
+        self.assertEqual({item["worker_id"] for item in items}, {"worker-1"})
+        for item, frame in zip(items, frames):
+            self.assertEqual(tuple(item["tensor"].shape), (1, 20, 30, 3))
+            self.assertEqual(item["tensor"].dtype, torch.float32)
+            expected = frame.astype(np.float32) / 255.0
+            self.assertTrue(np.array_equal(item["tensor"][0].numpy(), expected))
+
+    async def test_non_final_chunk_never_marks_worker_done(self):
+        response = await job_routes.job_complete_endpoint(
+            self._request_for(self._frames(count=2), "png", is_last=False)
+        )
+
+        self.assertEqual(response.status, 200)
+        self.assertFalse(any(item["is_last"] for item in await self._drain()))
+
+    async def test_audio_is_attached_only_to_final_frame(self):
+        audio = JobCompleteAudioPayloadTests._encoded_audio_payload(None)
+        response = await job_routes.job_complete_endpoint(
+            self._request_for(self._frames(count=3), "jpeg", audio=audio)
+        )
+
+        self.assertEqual(response.status, 200)
+        items = await self._drain()
+        self.assertEqual([item["audio"] is not None for item in items], [False, False, True])
+        self.assertEqual(tuple(items[-1]["audio"]["waveform"].shape), (1, 2, 4))
+
+    async def test_rejects_frame_count_mismatch(self):
+        response = await job_routes.job_complete_endpoint(
+            self._request_for(self._frames(count=2), "ffv1", count=5)
+        )
+
+        self.assertEqual(response.status, 400)
+        self.assertIn("frames", response.payload["message"])
+        self.assertTrue(self.queue.empty())
+
+    async def test_rejects_bad_envelope(self):
+        response = await job_routes.job_complete_endpoint(
+            self._request_for(self._frames(count=1), "png", format="gif", is_last="yes")
+        )
+
+        self.assertEqual(response.status, 400)
+        self.assertIn("format", response.payload["message"])
+        self.assertIn("is_last", response.payload["message"])
+
+    async def test_rejects_invalid_envelope_json(self):
+        response = await job_routes.job_complete_endpoint(_FakeMultipartRequest("{not json", [b"x"]))
+
+        self.assertEqual(response.status, 400)
+        self.assertIn("envelope", response.payload["message"])
+
+    async def test_unknown_job_returns_404(self):
+        job_routes.prompt_server.distributed_pending_jobs = {}
+        with patch.object(job_routes, "JOB_INIT_GRACE_PERIOD", 0.0):
+            response = await job_routes.job_complete_endpoint(self._request_for(self._frames(count=1), "png"))
+
+        self.assertEqual(response.status, 404)
+
+    async def test_legacy_json_payload_still_accepted(self):
+        request = _FakeRequest(
+            {
+                "job_id": "job-z",
+                "worker_id": "worker-1",
+                "batch_idx": 3,
+                "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, 200)
+        items = await self._drain()
+        self.assertEqual(len(items), 1)
+        self.assertEqual((items[0]["image_index"], items[0]["tensor"]), (3, "tensor-data"))
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/api/test_job_routes.py

@@ -1,360 +0,0 @@
-import importlib.util
-import sys
-import types
-import unittest
-import asyncio
-import base64
-from dataclasses import dataclass
-from pathlib import Path
-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):
-        self.payload = payload
-        self.status = status
-
-
-class _FakeRequest:
-    def __init__(self, payload, headers=None):
-        self._payload = payload
-        self.headers = headers or {}
-
-    async def json(self):
-        return self._payload
-
-
-def _load_job_routes_module():
-    module_path = Path(__file__).resolve().parents[2] / "api" / "job_routes.py"
-    package_name = "dist_api_queue_testpkg"
-
-    bootstrap_test_package(package_name, with_api=True, with_utils=True)
-
-    # aiohttp.web stub
-    created_aiohttp_stub = install_aiohttp_stub(
-        lambda payload, status=200: _FakeResponse(payload, status=status)
-    )
-
-    # server module stub with route decorators
-    prompt_server_instance = types.SimpleNamespace(
-        distributed_jobs_lock=None,
-        distributed_pending_jobs={},
-    )
-    install_server_stub(prompt_server_instance)
-
-    # torch stub (only needed to satisfy import)
-    created_torch_stub = False
-    if "torch" not in sys.modules:
-        created_torch_stub = True
-        torch_module = types.ModuleType("torch")
-        torch_module.cuda = types.SimpleNamespace(
-            is_available=lambda: False,
-            empty_cache=lambda: None,
-            ipc_collect=lambda: None,
-        )
-        sys.modules["torch"] = torch_module
-
-    # PIL stub (only needed to satisfy import)
-    created_pil_stub = False
-    if "PIL" not in sys.modules:
-        created_pil_stub = True
-        pil_module = types.ModuleType("PIL")
-        image_module = types.ModuleType("PIL.Image")
-        pil_module.Image = image_module
-        sys.modules["PIL"] = pil_module
-        sys.modules["PIL.Image"] = image_module
-
-    logging_module = types.ModuleType(f"{package_name}.utils.logging")
-    logging_module.debug_log = lambda *_args, **_kwargs: None
-    logging_module.log = lambda *_args, **_kwargs: None
-    sys.modules[f"{package_name}.utils.logging"] = logging_module
-
-    image_module = types.ModuleType(f"{package_name}.utils.image")
-    image_module.pil_to_tensor = lambda *_args, **_kwargs: None
-    image_module.ensure_contiguous = lambda tensor: tensor
-    sys.modules[f"{package_name}.utils.image"] = image_module
-
-    audio_payload_module = types.ModuleType(f"{package_name}.utils.audio_payload")
-
-    def _decode_audio_payload(payload):
-        if payload is None:
-            return None
-        if not isinstance(payload, dict):
-            raise ValueError("Field 'audio' must be an object when provided.")
-
-        encoded = payload.get("data")
-        shape = payload.get("shape")
-        dtype = payload.get("dtype", "float32")
-        sample_rate = payload.get("sample_rate", 44100)
-        if not isinstance(encoded, str) or not encoded.strip():
-            raise ValueError("Field 'audio.data' must be a non-empty base64 string.")
-        if not isinstance(shape, list) or len(shape) != 3:
-            raise ValueError("Field 'audio.shape' must be a 3-item list.")
-        if dtype != "float32":
-            raise ValueError("Field 'audio.dtype' must be 'float32'.")
-        try:
-            shape_tuple = tuple(int(dim) for dim in shape)
-        except Exception as exc:
-            raise ValueError("Field 'audio.shape' must contain integers.") from exc
-
-        raw = base64.b64decode(encoded, validate=True)
-        expected_bytes = int(np.prod(shape_tuple, dtype=np.int64)) * 4
-        if len(raw) != expected_bytes:
-            raise ValueError("Field 'audio.data' byte size mismatch.")
-
-        waveform = torch.from_numpy(np.frombuffer(raw, dtype=np.float32).reshape(shape_tuple).copy())
-        return {"waveform": waveform, "sample_rate": int(sample_rate)}
-
-    audio_payload_module.decode_audio_payload = _decode_audio_payload
-    sys.modules[f"{package_name}.utils.audio_payload"] = audio_payload_module
-
-    network_module = types.ModuleType(f"{package_name}.utils.network")
-
-    async def _handle_api_error(_request, error, status=500):
-        return _FakeResponse({"status": "error", "message": str(error)}, status=status)
-
-    network_module.handle_api_error = _handle_api_error
-    sys.modules[f"{package_name}.utils.network"] = network_module
-
-    constants_module = types.ModuleType(f"{package_name}.utils.constants")
-    constants_module.MEMORY_CLEAR_DELAY = 0.0
-    constants_module.JOB_INIT_GRACE_PERIOD = 10.0
-    sys.modules[f"{package_name}.utils.constants"] = constants_module
-
-    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
-
-    queue_orchestration_module = types.ModuleType(f"{package_name}.api.queue_orchestration")
-    queue_orchestration_module.orchestrate_distributed_execution = AsyncMock(return_value=("prompt_dist", 1))
-    sys.modules[f"{package_name}.api.queue_orchestration"] = queue_orchestration_module
-
-    @dataclass(frozen=True)
-    class _QueuePayload:
-        prompt: dict
-        workflow_meta: object
-        client_id: str
-        delegate_master: object
-        enabled_worker_ids: list
-        trace_execution_id: object
-
-    def _parse_queue_request_payload(data):
-        if not isinstance(data, dict):
-            raise ValueError("Expected a JSON object body")
-        prompt = data.get("prompt")
-        if not isinstance(prompt, dict):
-            raise ValueError("Field 'prompt' must be an object")
-        enabled = data.get("enabled_worker_ids")
-        if not isinstance(enabled, list):
-            raise ValueError("enabled_worker_ids required")
-        client_id = data.get("client_id")
-        if not isinstance(client_id, str) or not client_id.strip():
-            raise ValueError("client_id required")
-        return _QueuePayload(
-            prompt=prompt,
-            workflow_meta=data.get("workflow"),
-            client_id=client_id,
-            delegate_master=data.get("delegate_master"),
-            enabled_worker_ids=enabled,
-            trace_execution_id=data.get("trace_execution_id"),
-        )
-
-    queue_request_module = types.ModuleType(f"{package_name}.api.queue_request")
-    queue_request_module.parse_queue_request_payload = _parse_queue_request_payload
-    sys.modules[f"{package_name}.api.queue_request"] = queue_request_module
-
-    spec = importlib.util.spec_from_file_location(f"{package_name}.api.job_routes", module_path)
-    module = importlib.util.module_from_spec(spec)
-    assert spec is not None and spec.loader is not None
-    spec.loader.exec_module(module)
-
-    cleanup_optional_module("aiohttp", created_aiohttp_stub)
-    cleanup_optional_module("torch", created_torch_stub)
-    if created_pil_stub:
-        cleanup_optional_module("PIL.Image", True)
-        cleanup_optional_module("PIL", True)
-
-    return module
-
-
-job_routes = _load_job_routes_module()
-
-
-class DistributedQueueEndpointTests(unittest.IsolatedAsyncioTestCase):
-    async def test_distributed_queue_happy_path_returns_prompt_id(self):
-        request = _FakeRequest(
-            {
-                "prompt": {"1": {"class_type": "Node"}},
-                "enabled_worker_ids": ["w1"],
-                "client_id": "client-1",
-                "auto_prepare": True,
-            }
-        )
-        with patch.object(
-            job_routes,
-            "orchestrate_distributed_execution",
-            new=AsyncMock(return_value=("prompt_123", 2)),
-        ):
-            response = await job_routes.distributed_queue_endpoint(request)
-
-        self.assertEqual(response.status, 200)
-        self.assertEqual(response.payload.get("prompt_id"), "prompt_123")
-
-    async def test_distributed_queue_missing_prompt_returns_400(self):
-        request = _FakeRequest(
-            {
-                "enabled_worker_ids": ["w1"],
-                "client_id": "client-1",
-            }
-        )
-        response = await job_routes.distributed_queue_endpoint(request)
-        self.assertEqual(response.status, 400)
-        self.assertIn("prompt", response.payload.get("message", "").lower())
-
-    async def test_distributed_queue_missing_enabled_worker_ids_returns_400(self):
-        request = _FakeRequest(
-            {
-                "prompt": {"1": {"class_type": "Node"}},
-                "client_id": "client-1",
-            }
-        )
-        response = await job_routes.distributed_queue_endpoint(request)
-        self.assertEqual(response.status, 400)
-        self.assertIn("enabled_worker_ids", response.payload.get("message", "").lower())
-
-
-class JobCompleteAudioPayloadTests(unittest.IsolatedAsyncioTestCase):
-    def _encoded_audio_payload(self):
-        waveform = np.arange(8, dtype=np.float32).reshape(1, 2, 4)
-        return {
-            "sample_rate": 44100,
-            "shape": [1, 2, 4],
-            "dtype": "float32",
-            "data": base64.b64encode(waveform.tobytes()).decode("ascii"),
-        }
-
-    async def test_job_complete_accepts_audio_payload(self):
-        queue = asyncio.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",
-                "worker_id": "worker-1",
-                "batch_idx": 0,
-                "image": "data:image/png;base64,AAAA",
-                "audio": self._encoded_audio_payload(),
-                "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, 200)
-        queued = await queue.get()
-        self.assertEqual(queued["worker_id"], "worker-1")
-        self.assertTrue(queued["is_last"])
-        self.assertIsNotNone(queued["audio"])
-        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,
-            "shape": [1, 2],
-            "dtype": "float32",
-            "data": base64.b64encode(b"\x00\x00\x00\x00").decode("ascii"),
-        }
-        with self.assertRaises(ValueError):
-            job_routes._decode_audio_payload(bad)
-
-    def test_decode_audio_payload_rejects_bad_dtype(self):
-        payload = {
-            "sample_rate": 44100,
-            "shape": [1, 2, 4],
-            "dtype": "float16",
-            "data": base64.b64encode((np.zeros((1, 2, 4), dtype=np.float32)).tobytes()).decode("ascii"),
-        }
-        with self.assertRaises(ValueError):
-            job_routes._decode_audio_payload(payload)
-
-
-if __name__ == "__main__":
-    unittest.main()

tests/api/test_media_sync.py

@@ -4,29 +4,30 @@ 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"
 
-    bootstrap_test_package(
-        package_name,
-        with_api=True,
-        with_utils=True,
-        with_orchestration=True,
-    )
+    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
 
     logging_module = types.ModuleType(f"{package_name}.utils.logging")
     logging_module.debug_log = lambda *_args, **_kwargs: None
@@ -47,17 +48,22 @@ def _load_media_sync_module():
     trace_module.trace_info = lambda *_args, **_kwargs: None
     sys.modules[f"{package_name}.utils.trace_logger"] = trace_module
 
-    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
+    created_aiohttp_stub = False
+    if "aiohttp" not in sys.modules:
+        created_aiohttp_stub = True
+        aiohttp_module = types.ModuleType("aiohttp")
 
-    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 _ClientTimeout:
+            def __init__(self, total=None):
+                pass
 
-    created_aiohttp_stub = install_aiohttp_stub(
-        lambda payload, status=200: _AiohttpResponse(payload, status=status)
-    )
+        class _FormData:
+            def add_field(self, *args, **kwargs):
+                pass
+
+        aiohttp_module.ClientTimeout = _ClientTimeout
+        aiohttp_module.FormData = _FormData
+        sys.modules["aiohttp"] = aiohttp_module
 
     spec = importlib.util.spec_from_file_location(
         f"{package_name}.api.orchestration.media_sync",
@@ -67,7 +73,8 @@ def _load_media_sync_module():
     assert spec is not None and spec.loader is not None
     spec.loader.exec_module(module)
 
-    cleanup_optional_module("aiohttp", created_aiohttp_stub)
+    if created_aiohttp_stub:
+        sys.modules.pop("aiohttp", None)
 
     return module
 
@@ -157,26 +164,11 @@ class ConvertPathsForPlatformTests(unittest.TestCase):
 # ---------------------------------------------------------------------------
 
 class FindMediaReferencesTests(unittest.TestCase):
-    def test_finds_image_input(self):
-        prompt = {"1": {"class_type": "LoadImage", "inputs": {"image": "photo.png"}}}
-        refs = ms._find_media_references(prompt)
-        self.assertIn("photo.png", refs)
-
-    def test_finds_video_input(self):
-        prompt = {"1": {"class_type": "LoadVideo", "inputs": {"video": "clip.mp4"}}}
-        refs = ms._find_media_references(prompt)
-        self.assertIn("clip.mp4", refs)
-
     def test_finds_file_input_for_load_video(self):
         prompt = {"1": {"class_type": "LoadVideo", "inputs": {"file": "1 - Copy.mp4"}}}
         refs = ms._find_media_references(prompt)
         self.assertIn("1 - Copy.mp4", refs)
 
-    def test_finds_audio_input(self):
-        prompt = {"1": {"class_type": "LoadAudio", "inputs": {"audio": "track.wav"}}}
-        refs = ms._find_media_references(prompt)
-        self.assertIn("track.wav", refs)
-
     def test_strips_annotation_suffix(self):
         prompt = {"1": {"class_type": "LoadImage", "inputs": {"image": "photo.jpg [abc123]"}}}
         refs = ms._find_media_references(prompt)
@@ -253,7 +245,70 @@ class RewritePromptMediaInputsTests(unittest.TestCase):
 
 if __name__ == "__main__":
     unittest.main()
-    class _AiohttpResponse:
-        def __init__(self, payload, status=200):
-            self.payload = payload
-            self.status = status
+
+
+# ---------------------------------------------------------------------------
+# fetch_worker_has_node
+# ---------------------------------------------------------------------------
+
+class _FakeResponse:
+    def __init__(self, status, payload):
+        self.status = status
+        self._payload = payload
+
+    async def json(self):
+        if isinstance(self._payload, Exception):
+            raise self._payload
+        return self._payload
+
+    async def __aenter__(self):
+        return self
+
+    async def __aexit__(self, *exc):
+        return False
+
+
+class _FakeSession:
+    def __init__(self, response=None, error=None):
+        self.response = response
+        self.error = error
+        self.urls = []
+
+    def get(self, url, timeout=None):
+        self.urls.append(url)
+        if self.error:
+            raise self.error
+        return self.response
+
+
+class FetchWorkerHasNodeTests(unittest.TestCase):
+    def _probe(self, session):
+        import asyncio
+
+        async def _get_session():
+            return session
+
+        original = ms.get_client_session
+        ms.get_client_session = _get_session
+        try:
+            return asyncio.run(ms.fetch_worker_has_node({"id": "w1"}, "DistributedWorkerSink"))
+        finally:
+            ms.get_client_session = original
+
+    def test_registered_node_is_reported(self):
+        session = _FakeSession(_FakeResponse(200, {"DistributedWorkerSink": {"name": "DistributedWorkerSink"}}))
+        self.assertTrue(self._probe(session))
+        self.assertEqual(session.urls, ["http://localhost/object_info/DistributedWorkerSink"])
+
+    def test_unknown_node_returns_empty_object(self):
+        # ComfyUI answers 200 {} for classes that are not registered.
+        self.assertFalse(self._probe(_FakeSession(_FakeResponse(200, {}))))
+
+    def test_http_error_counts_as_missing(self):
+        self.assertFalse(self._probe(_FakeSession(_FakeResponse(404, None))))
+
+    def test_bad_json_counts_as_missing(self):
+        self.assertFalse(self._probe(_FakeSession(_FakeResponse(200, ValueError("bad json")))))
+
+    def test_connection_error_counts_as_missing(self):
+        self.assertFalse(self._probe(_FakeSession(error=OSError("unreachable"))))

tests/api/test_queue_orchestration.py

@@ -1,195 +0,0 @@
-import asyncio
-import importlib.util
-import sys
-import types
-import unittest
-from pathlib import Path
-
-from tests.api.harness import (
-    bootstrap_test_package,
-    install_server_stub,
-)
-
-
-def _bootstrap_package(package_name):
-    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)
-    install_server_stub(prompt_server_instance)
-
-    async_helpers_module = types.ModuleType(f"{package_name}.utils.async_helpers")
-
-    async def _queue_prompt_payload(prompt_obj, workflow_meta, client_id):
-        _ = (prompt_obj, workflow_meta, client_id)
-        return "prompt-id"
-
-    async_helpers_module.queue_prompt_payload = _queue_prompt_payload
-    sys.modules[f"{package_name}.utils.async_helpers"] = async_helpers_module
-
-    config_module = types.ModuleType(f"{package_name}.utils.config")
-    config_module.load_config = lambda: {"settings": {}, "workers": []}
-    sys.modules[f"{package_name}.utils.config"] = config_module
-
-    constants_module = types.ModuleType(f"{package_name}.utils.constants")
-    constants_module.ORCHESTRATION_MEDIA_SYNC_CONCURRENCY = 2
-    constants_module.ORCHESTRATION_MEDIA_SYNC_TIMEOUT = 120.0
-    constants_module.ORCHESTRATION_WORKER_PROBE_CONCURRENCY = 8
-    constants_module.ORCHESTRATION_WORKER_PREP_CONCURRENCY = 4
-    sys.modules[f"{package_name}.utils.constants"] = constants_module
-
-    logging_module = types.ModuleType(f"{package_name}.utils.logging")
-    logging_module.debug_log = lambda *_args, **_kwargs: None
-    logging_module.log = lambda *_args, **_kwargs: None
-    sys.modules[f"{package_name}.utils.logging"] = logging_module
-
-    network_module = types.ModuleType(f"{package_name}.utils.network")
-    network_module.build_master_url = lambda **_kwargs: "http://127.0.0.1:8188"
-    sys.modules[f"{package_name}.utils.network"] = network_module
-
-    trace_logger_module = types.ModuleType(f"{package_name}.utils.trace_logger")
-    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.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
-
-    dispatch_module = types.ModuleType(f"{package_name}.api.orchestration.dispatch")
-    dispatch_module.dispatch_worker_prompt = lambda *args, **kwargs: None
-    dispatch_module.rank_workers_by_load = lambda workers, **kwargs: workers
-    dispatch_module.select_active_workers = lambda workers, *_args, **_kwargs: (workers, False)
-    dispatch_module.select_least_busy_worker = lambda workers, **_kwargs: workers[0] if workers else None
-    sys.modules[f"{package_name}.api.orchestration.dispatch"] = dispatch_module
-
-    media_sync_module = types.ModuleType(f"{package_name}.api.orchestration.media_sync")
-    media_sync_module.convert_paths_for_platform = lambda prompt, _sep: prompt
-    media_sync_module.fetch_worker_path_separator = lambda *_args, **_kwargs: "/"
-    media_sync_module.sync_worker_media = lambda *_args, **_kwargs: None
-    sys.modules[f"{package_name}.api.orchestration.media_sync"] = media_sync_module
-
-    prompt_transform_module = types.ModuleType(f"{package_name}.api.orchestration.prompt_transform")
-    prompt_transform_module.PromptIndex = object
-    prompt_transform_module.apply_participant_overrides = lambda prompt, *_args, **_kwargs: prompt
-    prompt_transform_module.find_nodes_by_class = lambda _prompt, _class_name: []
-    prompt_transform_module.generate_job_id_map = lambda _index, _prefix: {}
-    prompt_transform_module.prepare_delegate_master_prompt = lambda prompt, _ids: prompt
-    prompt_transform_module.prune_prompt_for_worker = lambda prompt: prompt
-    sys.modules[f"{package_name}.api.orchestration.prompt_transform"] = prompt_transform_module
-
-
-def _load_module(package_name, module_rel_path, module_name):
-    module_path = Path(__file__).resolve().parents[2] / 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_queue_orchestration_module():
-    package_name = "dist_queue_orch_testpkg"
-    _bootstrap_package(package_name)
-    return _load_module(package_name, "api/queue_orchestration.py", "api.queue_orchestration")
-
-
-queue_orchestration = _load_queue_orchestration_module()
-
-
-class QueueOrchestrationHelpersTests(unittest.TestCase):
-    def test_resolve_enabled_workers_filters_disabled_and_requested(self):
-        config = {
-            "workers": [
-                {"id": "1", "enabled": True, "host": "a", "port": 8188},
-                {"id": "2", "enabled": False, "host": "b", "port": 8189},
-                {"id": "3", "enabled": True, "host": "c", "port": "bad-port"},
-            ]
-        }
-
-        enabled_only = queue_orchestration._resolve_enabled_workers(config)
-        requested = queue_orchestration._resolve_enabled_workers(config, requested_ids={"2", "3"})
-
-        self.assertEqual([worker["id"] for worker in enabled_only], ["1", "3"])
-        self.assertEqual([worker["id"] for worker in requested], ["2", "3"])
-        self.assertEqual(requested[1]["port"], 8188)
-
-    def test_resolve_orchestration_limits_uses_defaults_for_invalid_values(self):
-        config = {
-            "settings": {
-                "worker_probe_concurrency": "x",
-                "worker_prep_concurrency": "3",
-                "media_sync_concurrency": "-1",
-                "media_sync_timeout_seconds": "42.5",
-            }
-        }
-
-        limits = queue_orchestration._resolve_orchestration_limits(config)
-
-        self.assertEqual(limits[0], 8)   # default
-        self.assertEqual(limits[1], 3)   # parsed
-        self.assertEqual(limits[2], 2)   # default
-        self.assertEqual(limits[3], 42.5)
-
-    def test_is_load_balance_enabled_accepts_common_forms(self):
-        self.assertTrue(queue_orchestration._is_load_balance_enabled(True))
-        self.assertTrue(queue_orchestration._is_load_balance_enabled(1))
-        self.assertTrue(queue_orchestration._is_load_balance_enabled(" yes "))
-        self.assertFalse(queue_orchestration._is_load_balance_enabled(0))
-        self.assertFalse(queue_orchestration._is_load_balance_enabled("off"))
-
-    def test_prompt_requests_load_balance_detects_collector_flag(self):
-        prompt_index = types.SimpleNamespace(
-            nodes_for_class=lambda class_name: ["10"] if class_name == "DistributedCollector" else [],
-            inputs_by_node={"10": {"load_balance": "true"}},
-        )
-        self.assertTrue(queue_orchestration._prompt_requests_load_balance(prompt_index))
-
-    def test_ensure_distributed_state_initializes_attrs(self):
-        prompt_server = types.SimpleNamespace()
-        queue_orchestration.ensure_distributed_state(prompt_server)
-        self.assertTrue(hasattr(prompt_server, "distributed_pending_jobs"))
-        self.assertTrue(hasattr(prompt_server, "distributed_jobs_lock"))
-        self.assertIsInstance(prompt_server.distributed_jobs_lock, asyncio.Lock)
-
-
-if __name__ == "__main__":
-    unittest.main()

tests/api/test_transfer_integration.py

@@ -0,0 +1,219 @@
+"""End-to-end: real worker send path -> real aiohttp server -> job_complete endpoint."""
+import asyncio
+import importlib.util
+import sys
+from pathlib import Path
+
+import numpy as np
+import pytest
+import torch
+
+_TESTS_DIR = Path(__file__).resolve().parents[1]
+
+
+def _load_from(path, name):
+    spec = importlib.util.spec_from_file_location(name, path)
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    return module
+
+
+def _load_modules():
+    real_aiohttp = sys.modules["aiohttp"]  # imported inside the isolated test
+    queue_tests = _load_from(_TESTS_DIR / "api" / "test_distributed_queue.py", "dist_transfer_it_queue")
+    collector_tests = _load_from(_TESTS_DIR / "test_collector_list_inputs.py", "dist_transfer_it_collector")
+    job_routes = queue_tests._load_job_routes_module()
+    collector = collector_tests._load_collector_module()
+
+    # Error responses must be real HTTP responses here (the unit-test harness
+    # returns plain objects); mirror utils.network.handle_api_error.
+    from aiohttp import web
+
+    async def handle_api_error(_request, error, status=500):
+        if isinstance(error, list):
+            return web.json_response({"errors": [str(item) for item in error]}, status=status)
+        return web.json_response({"error": str(error)}, status=status)
+
+    job_routes.handle_api_error = handle_api_error
+    # The queue-test harness also fakes aiohttp.web for route modules; the
+    # endpoint must build real responses for the real server used here.
+    job_routes.web = web
+    # The collector harness stubs aiohttp; this test needs the real client.
+    sys.modules["aiohttp"] = real_aiohttp
+    collector.aiohttp = real_aiohttp
+    return job_routes, collector
+
+
+@pytest.fixture(autouse=True)
+def _isolate_sys_modules():
+    """The loaders below install stub modules (server, comfy, aiohttp...); undo that."""
+    saved = dict(sys.modules)
+    # Earlier test modules may have left a stub `aiohttp` behind; drop it so the
+    # real package is imported for this test (restored afterwards).
+    for name in [n for n in sys.modules if n == "aiohttp" or n.startswith("aiohttp.")]:
+        if getattr(sys.modules[name], "__file__", None) is None:
+            del sys.modules[name]
+    yield
+    sys.modules.clear()
+    sys.modules.update(saved)
+
+
+async def _run_transfer(job_routes, collector_module, images, fmt, *, client_max_size,
+                        request_limit, chunk_target, audio=None):
+    """Send `images` from a real worker node to the real completion route.
+
+    Returns (queued items, sizes of every request body the server received,
+    exception raised by the worker or None).
+    """
+    import aiohttp
+    from aiohttp import web
+
+    queue = asyncio.Queue()
+    job_routes.prompt_server.distributed_jobs_lock = asyncio.Lock()
+    job_routes.prompt_server.distributed_pending_jobs = {"job-it": queue}
+    body_sizes = []
+
+    @web.middleware
+    async def record_size(request, handler):
+        body_sizes.append(request.content_length)
+        return await handler(request)
+
+    app = web.Application(client_max_size=client_max_size, middlewares=[record_size])
+    app.router.add_post("/distributed/job_complete", job_routes.job_complete_endpoint)
+    runner = web.AppRunner(app)
+    await runner.setup()
+    site = web.TCPSite(runner, "127.0.0.1", 0)
+    await site.start()
+    port = site._server.sockets[0].getsockname()[1]
+
+    error = None
+    try:
+        async with aiohttp.ClientSession() as session:
+            async def _get_session():
+                return session
+
+            collector_module.get_client_session = _get_session
+            collector_module.TRANSFER_MAX_REQUEST_BYTES = request_limit
+            collector_module.TRANSFER_MAX_CHUNK_BYTES = chunk_target
+            node = collector_module.DistributedCollectorNode()
+            try:
+                await node.send_batch_to_master(
+                    images, audio, "job-it", f"http://127.0.0.1:{port}", "worker-1", fmt, 90
+                )
+            except Exception as exc:  # surfaced to the caller for assertions
+                error = exc
+    finally:
+        await runner.cleanup()
+
+    items = []
+    while not queue.empty():
+        items.append(queue.get_nowait())
+    return items, body_sizes, error
+
+
+def _noise_images(count, height=360, width=640, seed=97):
+    return torch.from_numpy(np.random.default_rng(seed).random((count, height, width, 3), dtype=np.float32))
+
+
+@pytest.mark.parametrize("fmt", ["h264", "ffv1"])
+def test_high_entropy_video_stays_under_server_upload_limit(fmt):
+    """Encoder flush used to push finalized chunks past the limit (HTTP 400)."""
+    pytest.importorskip("aiohttp")
+    job_routes, collector_module = _load_modules()
+    images = _noise_images(24)
+    limit = 2 * 1024 * 1024
+    reserve = collector_module.DistributedCollectorNode.TRANSFER_ENVELOPE_RESERVE_BYTES
+
+    items, body_sizes, error = asyncio.run(_run_transfer(
+        job_routes, collector_module, images, fmt,
+        client_max_size=limit, request_limit=limit, chunk_target=limit - reserve))
+
+    assert error is None, error
+    assert len(body_sizes) > 1 and max(body_sizes) <= limit
+    assert [item["image_index"] for item in items] == list(range(24))
+    assert [item["is_last"] for item in items] == [False] * 23 + [True]
+
+
+def test_server_rejection_for_size_is_reported_not_retried_as_legacy():
+    """A worker configured above the master's limit fails loudly instead of
+    silently resending the whole batch as uncompressed PNG."""
+    pytest.importorskip("aiohttp")
+    job_routes, collector_module = _load_modules()
+    images = _noise_images(8)
+
+    items, body_sizes, error = asyncio.run(_run_transfer(
+        job_routes, collector_module, images, "ffv1",
+        client_max_size=1024 * 1024, request_limit=8 * 1024 * 1024, chunk_target=4 * 1024 * 1024))
+
+    assert isinstance(error, RuntimeError)
+    assert "HTTP 413" in str(error) and "--max-upload-size" in str(error)
+    assert len(body_sizes) == 1
+    assert items == []
+
+
+def test_frame_that_cannot_fit_sends_nothing():
+    pytest.importorskip("aiohttp")
+    job_routes, collector_module = _load_modules()
+    images = _noise_images(3)
+    reserve = collector_module.DistributedCollectorNode.TRANSFER_ENVELOPE_RESERVE_BYTES
+
+    items, body_sizes, error = asyncio.run(_run_transfer(
+        job_routes, collector_module, images, "png",
+        client_max_size=100 * 1024 * 1024, request_limit=reserve + 64 * 1024, chunk_target=64 * 1024))
+
+    assert isinstance(error, collector_module.transfer_codec.FrameTooLargeError)
+    assert body_sizes == []
+    assert items == []
+
+
+@pytest.mark.parametrize("fmt", ["png", "ffv1", "h264", "jpeg"])
+def test_worker_to_master_over_real_http(fmt):
+    # Import here, not at module level: other test harnesses only install their
+    # aiohttp stubs when the real package has not been imported yet.
+    aiohttp = pytest.importorskip("aiohttp")
+    from aiohttp import web
+
+    job_routes, collector_module = _load_modules()
+    images = torch.rand(5, 24, 36, 3)
+
+    async def scenario():
+        queue = asyncio.Queue()
+        job_routes.prompt_server.distributed_jobs_lock = asyncio.Lock()
+        job_routes.prompt_server.distributed_pending_jobs = {"job-it": queue}
+
+        app = web.Application()
+        app.router.add_post("/distributed/job_complete", job_routes.job_complete_endpoint)
+        runner = web.AppRunner(app)
+        await runner.setup()
+        site = web.TCPSite(runner, "127.0.0.1", 0)
+        await site.start()
+        port = site._server.sockets[0].getsockname()[1]
+
+        async with aiohttp.ClientSession() as session:
+            async def _get_session():
+                return session
+
+            collector_module.get_client_session = _get_session
+            collector_module.TRANSFER_MAX_CHUNK_BYTES = 2000  # several requests for still formats
+            node = collector_module.DistributedCollectorNode()
+            await node.send_batch_to_master(
+                images, None, "job-it", f"http://127.0.0.1:{port}", "worker-1", fmt, 90
+            )
+        await runner.cleanup()
+
+        items = []
+        while not queue.empty():
+            items.append(queue.get_nowait())
+        return items
+
+    items = asyncio.run(scenario())
+
+    assert [item["image_index"] for item in items] == [0, 1, 2, 3, 4]
+    assert [item["is_last"] for item in items] == [False, False, False, False, True]
+    expected = np.clip(255.0 * images.numpy(), 0, 255).astype(np.uint8).astype(np.float32) / 255.0
+    got = torch.cat([item["tensor"] for item in items]).numpy()
+    assert got.shape == expected.shape
+    if fmt in ("png", "ffv1"):
+        assert np.array_equal(got, expected)
+    else:
+        assert np.abs(got - expected).mean() < 0.25  # random noise: only sanity-check lossy codecs

tests/api/test_tunnel_routes.py

@@ -1,95 +0,0 @@
-import importlib.util
-import sys
-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):
-        self.payload = payload
-        self.status = status
-
-
-def _load_tunnel_routes_module():
-    module_path = Path(__file__).resolve().parents[2] / "api" / "tunnel_routes.py"
-    package_name = "dist_tunnel_routes_testpkg"
-
-    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")
-
-    class _TunnelManager:
-        def get_status(self):
-            return {"status": "running"}
-
-        async def start_tunnel(self):
-            return {"status": "running", "public_url": "https://example.trycloudflare.com"}
-
-        async def stop_tunnel(self):
-            return {"status": "stopped"}
-
-    cloudflare_module.cloudflare_tunnel_manager = _TunnelManager()
-    sys.modules[f"{package_name}.utils.cloudflare"] = cloudflare_module
-
-    config_module = types.ModuleType(f"{package_name}.utils.config")
-    config_module.load_config = lambda: {"master": {"host": "https://master.example.com"}}
-    sys.modules[f"{package_name}.utils.config"] = config_module
-
-    logging_module = types.ModuleType(f"{package_name}.utils.logging")
-    logging_module.debug_log = lambda *_args, **_kwargs: None
-    logging_module.log = lambda *_args, **_kwargs: None
-    sys.modules[f"{package_name}.utils.logging"] = logging_module
-
-    network_module = types.ModuleType(f"{package_name}.utils.network")
-
-    async def _handle_api_error(_request, error, status=500):
-        return _FakeResponse({"status": "error", "message": str(error)}, status=status)
-
-    network_module.handle_api_error = _handle_api_error
-    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
-    sys.modules[spec.name] = module
-    spec.loader.exec_module(module)
-    return module
-
-
-tunnel_routes = _load_tunnel_routes_module()
-
-
-class TunnelRoutesTests(unittest.IsolatedAsyncioTestCase):
-    async def test_status_endpoint_returns_tunnel_and_master_host(self):
-        response = await tunnel_routes.tunnel_status_endpoint(request=None)
-        self.assertEqual(response.status, 200)
-        self.assertEqual(response.payload["status"], "success")
-        self.assertEqual(response.payload["tunnel"]["status"], "running")
-        self.assertEqual(response.payload["master_host"], "https://master.example.com")
-
-    async def test_start_and_stop_endpoints_return_success_payload(self):
-        start_response = await tunnel_routes.tunnel_start_endpoint(request=None)
-        stop_response = await tunnel_routes.tunnel_stop_endpoint(request=None)
-
-        self.assertEqual(start_response.status, 200)
-        self.assertEqual(start_response.payload["tunnel"]["status"], "running")
-        self.assertEqual(stop_response.status, 200)
-        self.assertEqual(stop_response.payload["tunnel"]["status"], "stopped")
-
-
-if __name__ == "__main__":
-    unittest.main()

tests/api/test_usdu_routes.py

@@ -9,14 +9,6 @@ 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):
@@ -38,30 +30,43 @@ 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"
 
-    bootstrap_test_package(package_name, with_api=True, with_utils=True, with_upscale=True)
+    for mod_name in list(sys.modules):
+        if mod_name == package_name or mod_name.startswith(f"{package_name}."):
+            del sys.modules[mod_name]
 
-    schemas_module = types.ModuleType(f"{package_name}.api.schemas")
+    root_pkg = types.ModuleType(package_name)
+    root_pkg.__path__ = []
+    sys.modules[package_name] = root_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.")
+    api_pkg = types.ModuleType(f"{package_name}.api")
+    api_pkg.__path__ = []
+    sys.modules[f"{package_name}.api"] = api_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
+    upscale_pkg = types.ModuleType(f"{package_name}.upscale")
+    upscale_pkg.__path__ = []
+    sys.modules[f"{package_name}.upscale"] = upscale_pkg
 
-    install_request_guards_stub(package_name)
+    utils_pkg = types.ModuleType(f"{package_name}.utils")
+    utils_pkg.__path__ = []
+    sys.modules[f"{package_name}.utils"] = utils_pkg
 
     prompt_server_holder = {
         "value": types.SimpleNamespace(
@@ -70,23 +75,23 @@ def _load_usdu_routes_module():
         )
     }
 
-    created_aiohttp_stub = install_aiohttp_stub(
-        lambda payload, status=200: _FakeResponse(payload, status=status)
-    )
-    install_server_stub()
+    created_aiohttp_stub = False
+    if "aiohttp" not in sys.modules:
+        created_aiohttp_stub = True
+        aiohttp_module = types.ModuleType("aiohttp")
+        aiohttp_module.web = types.SimpleNamespace(
+            json_response=lambda payload, status=200: _FakeResponse(payload, status=status)
+        )
+        sys.modules["aiohttp"] = aiohttp_module
+
+    server_module = types.ModuleType("server")
+    server_module.PromptServer = types.SimpleNamespace(instance=types.SimpleNamespace(routes=_Routes()))
+    sys.modules["server"] = server_module
 
     logging_module = types.ModuleType(f"{package_name}.utils.logging")
     logging_module.debug_log = lambda *_args, **_kwargs: None
     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):
@@ -98,21 +103,6 @@ def _load_usdu_routes_module():
     job_store_module = types.ModuleType(f"{package_name}.upscale.job_store")
     job_store_module.MAX_PAYLOAD_SIZE = 1024
     job_store_module.ensure_tile_jobs_initialized = lambda: 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)
-            job_data.worker_status = {str(worker_id): 0 for worker_id in (enabled_workers or [])}
-            job_data.assigned_to_workers = {str(worker_id): [] for worker_id in (enabled_workers or [])}
-            indices = all_indices if all_indices is not None else list(range(int(batch_size or 0)))
-            for idx in indices:
-                await job_data.pending_images.put(int(idx))
-            prompt_server.distributed_pending_tile_jobs[multi_job_id] = job_data
-
-    job_store_module.init_dynamic_job = _init_dynamic_job
     sys.modules[f"{package_name}.upscale.job_store"] = job_store_module
 
     job_models_module = types.ModuleType(f"{package_name}.upscale.job_models")
@@ -160,7 +150,6 @@ def _load_usdu_routes_module():
     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: []
     parsers_module._parse_tiles_from_form = lambda _data: []
     sys.modules[f"{package_name}.upscale.payload_parsers"] = parsers_module
 
@@ -169,7 +158,8 @@ def _load_usdu_routes_module():
     assert spec is not None and spec.loader is not None
     spec.loader.exec_module(module)
 
-    cleanup_optional_module("aiohttp", created_aiohttp_stub)
+    if created_aiohttp_stub:
+        sys.modules.pop("aiohttp", None)
 
     module.web = types.SimpleNamespace(
         json_response=lambda payload, status=200: _FakeResponse(payload, status=status)
@@ -232,8 +222,7 @@ class USDURoutesTests(unittest.IsolatedAsyncioTestCase):
         response = await usdu_routes.request_image_endpoint(request)
 
         self.assertEqual(response.status, 200)
-        self.assertEqual(response.payload.get("kind"), "image")
-        self.assertEqual(response.payload.get("task_idx"), 7)
+        self.assertEqual(response.payload.get("image_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)
@@ -249,8 +238,7 @@ class USDURoutesTests(unittest.IsolatedAsyncioTestCase):
         response = await usdu_routes.request_image_endpoint(request)
 
         self.assertEqual(response.status, 200)
-        self.assertEqual(response.payload.get("kind"), "tile")
-        self.assertEqual(response.payload.get("task_idx"), 4)
+        self.assertEqual(response.payload.get("tile_idx"), 4)
         self.assertTrue(response.payload.get("batched_static"))
         self.assertEqual(job_data.assigned_to_workers["worker-a"], [4])
 

tests/api/test_worker_routes.py

@@ -8,14 +8,6 @@ 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):
@@ -24,11 +16,10 @@ class _FakeResponse:
 
 
 class _FakeRequest:
-    def __init__(self, payload=None, match_info=None, query=None, headers=None):
+    def __init__(self, payload=None, match_info=None, query=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
@@ -58,8 +49,8 @@ class _FakeHTTPClientSession:
         self._status = status
         self.calls = []
 
-    def get(self, url, params=None, headers=None, timeout=None):
-        self.calls.append({"url": url, "params": params, "headers": headers, "timeout": timeout})
+    def get(self, url, params=None, timeout=None):
+        self.calls.append({"url": url, "params": params, "timeout": timeout})
         return _FakeHTTPClientResponse(self._payload, status=self._status)
 
 
@@ -71,12 +62,12 @@ class _DummyWorkerManager:
         worker_id = str(worker["id"])
         self.processes[worker_id] = {
             "pid": 12345,
-            "log_file": f"/tmp/distributed_worker_{worker_id}.log",  # nosec B108 - deterministic fake test path
+            "log_file": f"/tmp/distributed_worker_{worker_id}.log",
             "process": None,
         }
         return 12345
 
-    def is_process_running(self, _pid):
+    def _is_process_running(self, _pid):
         return False
 
     def save_processes(self):
@@ -98,10 +89,24 @@ def _load_worker_routes_module():
     module_path = Path(__file__).resolve().parents[2] / "api" / "worker_routes.py"
     package_name = "dist_api_worker_testpkg"
 
-    bootstrap_test_package(package_name, with_api=True, with_utils=True, with_workers=True)
+    for mod_name in list(sys.modules):
+        if mod_name == package_name or mod_name.startswith(f"{package_name}."):
+            del sys.modules[mod_name]
+
+    root_pkg = types.ModuleType(package_name)
+    root_pkg.__path__ = []
+    sys.modules[package_name] = root_pkg
+
+    api_pkg = types.ModuleType(f"{package_name}.api")
+    api_pkg.__path__ = []
+    sys.modules[f"{package_name}.api"] = api_pkg
+
+    utils_pkg = types.ModuleType(f"{package_name}.utils")
+    utils_pkg.__path__ = []
+    sys.modules[f"{package_name}.utils"] = utils_pkg
 
     workers_pkg = types.ModuleType(f"{package_name}.workers")
-    workers_pkg.__path__ = []
+    workers_pkg.__path__ = [str(module_path.parents[1] / "workers")]
     workers_pkg.get_worker_manager = lambda: _DummyWorkerManager()
     sys.modules[f"{package_name}.workers"] = workers_pkg
 
@@ -114,10 +119,59 @@ 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 = install_aiohttp_stub(
-        lambda payload, status=200: _FakeResponse(payload, status=status)
-    )
-    install_server_stub()
+    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_torch_stub = False
     if "torch" not in sys.modules:
@@ -159,6 +213,7 @@ def _load_worker_routes_module():
         raise RuntimeError("not used in these tests")
 
     network_module.get_client_session = _get_client_session
+    network_module.get_server_port = lambda: 8189
     sys.modules[f"{package_name}.utils.network"] = network_module
 
     constants_module = types.ModuleType(f"{package_name}.utils.constants")
@@ -194,29 +249,21 @@ 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)
+    # Keep this module's response double independent of aiohttp import order.
+    module.web = types.SimpleNamespace(json_response=lambda payload, status=200: _FakeResponse(payload, status))
 
-    cleanup_optional_module("aiohttp", created_aiohttp_stub)
-    cleanup_optional_module("torch", created_torch_stub)
+    if created_aiohttp_stub:
+        sys.modules.pop("aiohttp", None)
+    if created_torch_stub:
+        sys.modules.pop("torch", None)
 
     return module
 
@@ -225,6 +272,38 @@ worker_routes = _load_worker_routes_module()
 
 
 class WorkerRoutesTests(unittest.IsolatedAsyncioTestCase):
+    async def test_network_info_reports_actual_master_and_available_ports(self):
+        with patch.object(worker_routes, "_get_cuda_info", return_value=(0, 1, 4)), \
+             patch.object(worker_routes, "get_network_ips", return_value=["127.0.0.1"]), \
+             patch.object(worker_routes, "load_config", return_value={"workers": []}), \
+             patch.object(worker_routes, "get_server_port", return_value=8189), \
+             patch.object(worker_routes, "allocate_worker_ports", return_value=[8191, 8192, 8193]) as allocate:
+            response = await worker_routes.get_network_info_endpoint(_FakeRequest())
+        self.assertEqual(response.payload["master_port"], 8189)
+        self.assertEqual(response.payload["local_worker_ports"], [8191, 8192, 8193])
+        self.assertEqual(response.payload["cuda_device_count"], 4)
+        allocate.assert_called_once_with(8189, [], 3)
+
+    async def test_network_info_does_not_reallocate_existing_configuration(self):
+        with patch.object(worker_routes, "_get_cuda_info", return_value=(0, 4, 4)), \
+             patch.object(worker_routes, "get_network_ips", return_value=[]), \
+             patch.object(worker_routes, "load_config", return_value={"workers": [{"id": "existing", "port": 8189}]}), \
+             patch.object(worker_routes, "allocate_worker_ports") as allocate:
+            response = await worker_routes.get_network_info_endpoint(_FakeRequest())
+        self.assertEqual(response.payload["local_worker_ports"], [])
+        allocate.assert_not_called()
+
+    async def test_launch_conflict_is_a_clear_client_error(self):
+        manager = _DummyWorkerManager()
+        config = {"workers": [{"id": "worker-a", "name": "Worker A", "port": 8189}]}
+        with patch.object(worker_routes, "get_worker_manager", return_value=manager), \
+             patch.object(worker_routes, "load_config", return_value=config), \
+             patch.object(manager, "launch_worker", side_effect=ValueError("Worker port 8189 conflicts with the master port 8189")):
+            response = await worker_routes.launch_worker_endpoint(_FakeRequest({"worker_id": "worker-a"}))
+        self.assertEqual(response.status, 400)
+        self.assertIn("master port 8189", response.payload["message"])
+        self.assertEqual(manager.processes, {})
+
     async def test_launch_worker_valid_id_returns_200(self):
         manager = _DummyWorkerManager()
         config = {"workers": [{"id": "worker-a", "name": "Worker A", "port": 8188}]}

tests/conftest.py

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

tests/fixtures/v1_node_contracts.json

@@ -0,0 +1,691 @@
+{
+  "DistributedCollector": {
+    "input": {
+      "required": {
+        "load_balance": [
+          "BOOLEAN",
+          {
+            "default": false,
+            "tooltip": "Run this workflow on one least-busy participant (master included when participating)."
+          }
+        ]
+      },
+      "optional": {
+        "images": [
+          "IMAGE"
+        ],
+        "audio": [
+          "AUDIO"
+        ]
+      },
+      "hidden": {
+        "multi_job_id": [
+          "STRING",
+          {
+            "default": ""
+          }
+        ],
+        "is_worker": [
+          "BOOLEAN",
+          {
+            "default": false
+          }
+        ],
+        "master_url": [
+          "STRING",
+          {
+            "default": ""
+          }
+        ],
+        "enabled_worker_ids": [
+          "STRING",
+          {
+            "default": "[]"
+          }
+        ],
+        "worker_batch_size": [
+          "INT",
+          {
+            "default": 1,
+            "min": 1,
+            "max": 1024
+          }
+        ],
+        "worker_id": [
+          "STRING",
+          {
+            "default": ""
+          }
+        ],
+        "pass_through": [
+          "BOOLEAN",
+          {
+            "default": false
+          }
+        ],
+        "delegate_only": [
+          "BOOLEAN",
+          {
+            "default": false
+          }
+        ]
+      }
+    },
+    "input_order": {
+      "required": [
+        "load_balance"
+      ],
+      "optional": [
+        "images",
+        "audio"
+      ],
+      "hidden": [
+        "multi_job_id",
+        "is_worker",
+        "master_url",
+        "enabled_worker_ids",
+        "worker_batch_size",
+        "worker_id",
+        "pass_through",
+        "delegate_only"
+      ]
+    },
+    "output": [
+      "IMAGE",
+      "AUDIO"
+    ],
+    "output_name": [
+      "images",
+      "audio"
+    ],
+    "output_is_list": [
+      false,
+      false
+    ],
+    "is_input_list": true,
+    "output_node": false,
+    "category": "image",
+    "display_name": "Distributed Collector"
+  },
+  "DistributedSeed": {
+    "input": {
+      "required": {
+        "seed": [
+          "INT",
+          {
+            "default": 1125899906842,
+            "min": 0,
+            "max": 1125899906842624,
+            "forceInput": false
+          }
+        ]
+      },
+      "hidden": {
+        "is_worker": [
+          "BOOLEAN",
+          {
+            "default": false
+          }
+        ],
+        "worker_id": [
+          "STRING",
+          {
+            "default": ""
+          }
+        ]
+      }
+    },
+    "input_order": {
+      "required": [
+        "seed"
+      ],
+      "hidden": [
+        "is_worker",
+        "worker_id"
+      ]
+    },
+    "output": [
+      "INT"
+    ],
+    "output_name": [
+      "seed"
+    ],
+    "output_is_list": [
+      false
+    ],
+    "is_input_list": false,
+    "output_node": false,
+    "category": "utils",
+    "display_name": "Distributed Seed"
+  },
+  "DistributedModelName": {
+    "input": {
+      "required": {
+        "text": [
+          "STRING",
+          {
+            "default": ""
+          }
+        ]
+      },
+      "hidden": {
+        "unique_id": "UNIQUE_ID",
+        "extra_pnginfo": "EXTRA_PNGINFO"
+      }
+    },
+    "input_order": {
+      "required": [
+        "text"
+      ],
+      "hidden": [
+        "unique_id",
+        "extra_pnginfo"
+      ]
+    },
+    "output": [
+      "*"
+    ],
+    "output_name": [
+      "output"
+    ],
+    "output_is_list": [
+      false
+    ],
+    "is_input_list": false,
+    "output_node": true,
+    "category": "utils",
+    "display_name": "Distributed Model Name"
+  },
+  "DistributedValue": {
+    "input": {
+      "required": {
+        "default_value": [
+          "STRING",
+          {
+            "default": ""
+          }
+        ],
+        "worker_values": [
+          "STRING",
+          {
+            "default": "{}"
+          }
+        ]
+      },
+      "hidden": {
+        "is_worker": [
+          "BOOLEAN",
+          {
+            "default": false
+          }
+        ],
+        "worker_id": [
+          "STRING",
+          {
+            "default": ""
+          }
+        ]
+      }
+    },
+    "input_order": {
+      "required": [
+        "default_value",
+        "worker_values"
+      ],
+      "hidden": [
+        "is_worker",
+        "worker_id"
+      ]
+    },
+    "output": [
+      "*"
+    ],
+    "output_name": [
+      "value"
+    ],
+    "output_is_list": [
+      false
+    ],
+    "is_input_list": false,
+    "output_node": false,
+    "category": "utils",
+    "display_name": "Distributed Value"
+  },
+  "ImageBatchDivider": {
+    "input": {
+      "required": {
+        "images": [
+          "IMAGE"
+        ],
+        "divide_by": [
+          "INT",
+          {
+            "default": 2,
+            "min": 1,
+            "max": 10,
+            "step": 1,
+            "display": "number",
+            "tooltip": "Number of parts to divide the batch into"
+          }
+        ]
+      }
+    },
+    "input_order": {
+      "required": [
+        "images",
+        "divide_by"
+      ]
+    },
+    "output": [
+      "*",
+      "*",
+      "*",
+      "*",
+      "*",
+      "*",
+      "*",
+      "*",
+      "*",
+      "*"
+    ],
+    "output_name": [
+      "batch_1",
+      "batch_2",
+      "batch_3",
+      "batch_4",
+      "batch_5",
+      "batch_6",
+      "batch_7",
+      "batch_8",
+      "batch_9",
+      "batch_10"
+    ],
+    "output_is_list": [
+      false,
+      false,
+      false,
+      false,
+      false,
+      false,
+      false,
+      false,
+      false,
+      false
+    ],
+    "is_input_list": false,
+    "output_node": true,
+    "category": "image",
+    "display_name": "Image Batch Divider"
+  },
+  "AudioBatchDivider": {
+    "input": {
+      "required": {
+        "audio": [
+          "AUDIO"
+        ],
+        "divide_by": [
+          "INT",
+          {
+            "default": 2,
+            "min": 1,
+            "max": 10,
+            "step": 1,
+            "display": "number",
+            "tooltip": "Number of sequential time segments to create"
+          }
+        ]
+      }
+    },
+    "input_order": {
+      "required": [
+        "audio",
+        "divide_by"
+      ]
+    },
+    "output": [
+      "*",
+      "*",
+      "*",
+      "*",
+      "*",
+      "*",
+      "*",
+      "*",
+      "*",
+      "*"
+    ],
+    "output_name": [
+      "audio_1",
+      "audio_2",
+      "audio_3",
+      "audio_4",
+      "audio_5",
+      "audio_6",
+      "audio_7",
+      "audio_8",
+      "audio_9",
+      "audio_10"
+    ],
+    "output_is_list": [
+      false,
+      false,
+      false,
+      false,
+      false,
+      false,
+      false,
+      false,
+      false,
+      false
+    ],
+    "is_input_list": false,
+    "output_node": true,
+    "category": "audio",
+    "display_name": "Audio Segment Divider"
+  },
+  "DistributedEmptyImage": {
+    "input": {
+      "required": {
+        "height": [
+          "INT",
+          {
+            "default": 64,
+            "min": 1,
+            "max": 4096,
+            "step": 1
+          }
+        ],
+        "width": [
+          "INT",
+          {
+            "default": 64,
+            "min": 1,
+            "max": 4096,
+            "step": 1
+          }
+        ],
+        "channels": [
+          "INT",
+          {
+            "default": 3,
+            "min": 1,
+            "max": 4,
+            "step": 1
+          }
+        ]
+      }
+    },
+    "input_order": {
+      "required": [
+        "height",
+        "width",
+        "channels"
+      ]
+    },
+    "output": [
+      "IMAGE"
+    ],
+    "output_name": [
+      "IMAGE"
+    ],
+    "output_is_list": [
+      false
+    ],
+    "is_input_list": false,
+    "output_node": false,
+    "category": "image",
+    "display_name": "Distributed Empty Image"
+  },
+  "UltimateSDUpscaleDistributed": {
+    "input": {
+      "required": {
+        "upscaled_image": [
+          "IMAGE"
+        ],
+        "model": [
+          "MODEL"
+        ],
+        "positive": [
+          "CONDITIONING"
+        ],
+        "negative": [
+          "CONDITIONING"
+        ],
+        "vae": [
+          "VAE"
+        ],
+        "seed": [
+          "INT",
+          {
+            "default": 0,
+            "min": 0,
+            "max": 18446744073709551615
+          }
+        ],
+        "steps": [
+          "INT",
+          {
+            "default": 20,
+            "min": 1,
+            "max": 10000
+          }
+        ],
+        "cfg": [
+          "FLOAT",
+          {
+            "default": 8.0,
+            "min": 0.0,
+            "max": 100.0
+          }
+        ],
+        "sampler_name": [
+          [
+            "euler",
+            "euler_cfg_pp",
+            "euler_ancestral",
+            "euler_ancestral_cfg_pp",
+            "heun",
+            "heunpp2",
+            "exp_heun_2_x0",
+            "exp_heun_2_x0_sde",
+            "dpm_2",
+            "dpm_2_ancestral",
+            "lms",
+            "dpm_fast",
+            "dpm_adaptive",
+            "dpmpp_2s_ancestral",
+            "dpmpp_2s_ancestral_cfg_pp",
+            "dpmpp_sde",
+            "dpmpp_sde_gpu",
+            "dpmpp_2m",
+            "dpmpp_2m_cfg_pp",
+            "dpmpp_2m_sde",
+            "dpmpp_2m_sde_gpu",
+            "dpmpp_2m_sde_heun",
+            "dpmpp_2m_sde_heun_gpu",
+            "dpmpp_3m_sde",
+            "dpmpp_3m_sde_gpu",
+            "ddpm",
+            "lcm",
+            "ipndm",
+            "ipndm_v",
+            "deis",
+            "cfgpp_ud10_ab",
+            "res_multistep",
+            "res_multistep_cfg_pp",
+            "res_multistep_ancestral",
+            "res_multistep_ancestral_cfg_pp",
+            "gradient_estimation",
+            "gradient_estimation_cfg_pp",
+            "er_sde",
+            "seeds_2",
+            "seeds_3",
+            "sa_solver",
+            "sa_solver_pece",
+            "ddim",
+            "uni_pc",
+            "uni_pc_bh2"
+          ]
+        ],
+        "scheduler": [
+          [
+            "simple",
+            "sgm_uniform",
+            "karras",
+            "exponential",
+            "ddim_uniform",
+            "beta",
+            "normal",
+            "linear_quadratic",
+            "kl_optimal"
+          ]
+        ],
+        "denoise": [
+          "FLOAT",
+          {
+            "default": 0.5,
+            "min": 0.0,
+            "max": 1.0,
+            "step": 0.01
+          }
+        ],
+        "tile_width": [
+          "INT",
+          {
+            "default": 512,
+            "min": 64,
+            "max": 2048,
+            "step": 8
+          }
+        ],
+        "tile_height": [
+          "INT",
+          {
+            "default": 512,
+            "min": 64,
+            "max": 2048,
+            "step": 8
+          }
+        ],
+        "padding": [
+          "INT",
+          {
+            "default": 32,
+            "min": 0,
+            "max": 256,
+            "step": 8
+          }
+        ],
+        "mask_blur": [
+          "INT",
+          {
+            "default": 8,
+            "min": 0,
+            "max": 256
+          }
+        ],
+        "force_uniform_tiles": [
+          "BOOLEAN",
+          {
+            "default": true
+          }
+        ],
+        "tiled_decode": [
+          "BOOLEAN",
+          {
+            "default": false
+          }
+        ]
+      },
+      "hidden": {
+        "multi_job_id": [
+          "STRING",
+          {
+            "default": ""
+          }
+        ],
+        "is_worker": [
+          "BOOLEAN",
+          {
+            "default": false
+          }
+        ],
+        "master_url": [
+          "STRING",
+          {
+            "default": ""
+          }
+        ],
+        "enabled_worker_ids": [
+          "STRING",
+          {
+            "default": "[]"
+          }
+        ],
+        "worker_id": [
+          "STRING",
+          {
+            "default": ""
+          }
+        ],
+        "tile_indices": [
+          "STRING",
+          {
+            "default": ""
+          }
+        ],
+        "dynamic_threshold": [
+          "INT",
+          {
+            "default": 8,
+            "min": 1,
+            "max": 64
+          }
+        ]
+      }
+    },
+    "input_order": {
+      "required": [
+        "upscaled_image",
+        "model",
+        "positive",
+        "negative",
+        "vae",
+        "seed",
+        "steps",
+        "cfg",
+        "sampler_name",
+        "scheduler",
+        "denoise",
+        "tile_width",
+        "tile_height",
+        "padding",
+        "mask_blur",
+        "force_uniform_tiles",
+        "tiled_decode"
+      ],
+      "hidden": [
+        "multi_job_id",
+        "is_worker",
+        "master_url",
+        "enabled_worker_ids",
+        "worker_id",
+        "tile_indices",
+        "dynamic_threshold"
+      ]
+    },
+    "output": [
+      "IMAGE"
+    ],
+    "output_name": [
+      "IMAGE"
+    ],
+    "output_is_list": [
+      false
+    ],
+    "is_input_list": false,
+    "output_node": false,
+    "category": "image/upscaling",
+    "display_name": "Ultimate SD Upscale Distributed (No Upscale)"
+  }
+}

tests/helpers/v3_runtime_check.py

@@ -0,0 +1,324 @@
+"""Exercise the actual loader, schemas and executor in a fresh CPU process.
+
+No HTTP listener, workers, model downloads or live installation changes.
+The original contracts were captured from 32ac027 using the same core.
+"""
+import asyncio
+import inspect
+import json
+import os
+from pathlib import Path
+import sys
+import tempfile
+from unittest.mock import patch
+
+from PIL import Image
+
+ROOT = Path(__file__).resolve().parents[2]
+COMFY_ROOT = Path(sys.argv[1]).resolve()
+sys.path.insert(0, str(COMFY_ROOT))
+os.environ['COMFYUI_IS_WORKER'] = '1'
+import comfy.cli_args
+comfy.cli_args.args.cpu = True
+comfy.cli_args.args.disable_assets = True
+from app.assets.manager import default_asset_manager
+from comfy_api.v0_0_2 import io
+import execution
+import nodes
+import server
+import torch
+
+BASELINE = json.loads((ROOT / 'tests/fixtures/v1_node_contracts.json').read_text())
+
+# Inputs intentionally added after the V1 baseline was captured. They must be
+# optional, appended after the baseline inputs, and keep legacy behaviour as
+# their default so saved workflows are unaffected.
+ADDED_OPTIONAL_INPUTS = {
+    'DistributedCollector': {
+        'transfer_format': [
+            ['legacy_png', 'png', 'ffv1', 'jpeg', 'webp', 'h264', 'h265', 'av1'],
+            {'default': 'legacy_png'},
+        ],
+        'transfer_quality': ['INT', {'default': 90, 'min': 1, 'max': 100, 'step': 1}],
+    },
+}
+
+# Nodes added after the V1 baseline; not part of saved-workflow contracts.
+# DistributedWorkerSink is internal (dev-only) and only inserted by the master.
+ADDED_NODES = {'DistributedWorkerSink'}
+
+
+def normalized_input(value):
+    kind = value[0]
+    opts = dict(value[1]) if len(value) > 1 else {}
+    if kind == 'STRING':
+        # V3 serializes the same default single-line widget explicitly.
+        opts.setdefault('multiline', False)
+    if kind == 'COMBO':
+        kind = opts['options']
+        opts = {key: val for key, val in opts.items() if key != 'options'}
+    if isinstance(kind, list):
+        # Single-selection is the V1 dropdown default as well.
+        opts.setdefault('multiselect', False)
+    return [kind, opts]
+
+
+async def map_node(cls, values, extra=None):
+    """Use the actual input parser and sanitized V3 executor clones."""
+    prepared, missing, hidden = execution.get_input_data(values, cls, unique_id='probe', extra_data=extra or {})
+    assert not missing, missing
+    # Use a worker thread for synchronous nodes: collector/upscale bridge back
+    # to PromptServer's running loop, just as ComfyUI's prompt worker does.
+    def run():
+        return asyncio.run(execution._async_map_node_over_list(
+            'v3-acceptance', 'probe', cls, prepared, cls.FUNCTION, v3_data=hidden))
+    return await asyncio.to_thread(run)
+
+
+def check_saved_workflows(mapping):
+    checked = 0
+    occurrences = 0
+    for path in sorted((ROOT / 'workflows').glob('*.json')):
+        workflow = json.loads(path.read_text())
+        local_nodes = {node['id']: node for node in workflow['nodes']}
+        for saved in workflow['nodes']:
+            if saved['type'] not in mapping:
+                continue
+            occurrences += 1
+            info = mapping[saved['type']].GET_NODE_INFO_V1()
+            inputs = {**info['input'].get('required', {}), **info['input'].get('optional', {})}
+            for socket in saved.get('inputs', []):
+                assert socket['name'] in inputs, (path.name, saved['id'], socket)
+                assert socket['type'] == inputs[socket['name']][0], (path.name, socket)
+            for index, output in enumerate(saved.get('outputs', [])):
+                assert output['type'] == info['output'][index], (path.name, index)
+            for link in workflow['links']:
+                if link[1] == saved['id']:
+                    assert link[2] < len(info['output']), (path.name, link)
+                    target = local_nodes[link[3]]['inputs'][link[4]]
+                    assert target['type'] == info['output'][link[2]], (path.name, link)
+        checked += 1
+    assert checked == 5 and occurrences == 10, (checked, occurrences)
+    print('SAVED_WORKFLOW_CONTRACTS_OK', checked, occurrences)
+
+
+async def check_execution(mapping, module, prompt_server, asset_manager):
+    result = await map_node(mapping['DistributedSeed'],
+                            {'seed': 123, 'is_worker': True, 'worker_id': 'worker_2'})
+    assert result[0].result == (126,)
+    values = json.dumps({'_type': 'INT', '3': '17'})
+    result = await map_node(mapping['DistributedValue'],
+                            {'default_value': '4', 'worker_values': values,
+                             'is_worker': True, 'worker_id': 'worker_2'})
+    assert result[0].result == (17,)
+    metadata = {'workflow': {'nodes': [{'id': 'probe', 'widgets_values': []}]}}
+    result = await map_node(mapping['DistributedModelName'], {'text': 'model.ckpt'},
+                            {'extra_pnginfo': metadata})
+    assert result[0].result == ('model.ckpt',)
+    assert result[0].ui == {'text': ['model.ckpt']}
+    assert metadata['workflow']['nodes'][0]['widgets_values'] == [['model.ckpt']]
+    result = await map_node(mapping['DistributedEmptyImage'],
+                            {'height': 8, 'width': 8, 'channels': 3})
+    empty = result[0].result[0]
+    assert empty.shape == (0, 8, 8, 3) and empty.numel() == 0
+    images = torch.arange(10 * 8 * 8 * 3, dtype=torch.float32).reshape(10, 8, 8, 3)
+    for count in (1, 3, 10):
+        result = await map_node(mapping['ImageBatchDivider'], {'images': images, 'divide_by': count})
+        assert len(result[0].result) == 10
+        assert torch.equal(torch.cat(result[0].result[:count]), images)
+    audio = {'waveform': torch.arange(33, dtype=torch.float32).reshape(1, 1, 33), 'sample_rate': 24000}
+    result = await map_node(mapping['AudioBatchDivider'], {'audio': audio, 'divide_by': 3})
+    assert len(result[0].result) == 10
+    assert torch.equal(torch.cat([item['waveform'] for item in result[0].result[:3]], dim=-1), audio['waveform'])
+    # INPUT_IS_LIST must preserve all images/audio but unwrap transport scalars.
+    collector = mapping['DistributedCollector']
+    prepared, missing, hidden = execution.get_input_data(
+        {'load_balance': False, 'multi_job_id': '', 'pass_through': True}, collector, 'collector-list')
+    assert not missing
+    prepared['images'] = [images[:2], images[2:]]
+    prepared['audio'] = [audio, audio]
+    result = await execution._async_map_node_over_list(
+        'v3-acceptance', 'collector-list', collector, prepared, collector.FUNCTION, v3_data=hidden)
+    assert torch.equal(result[0].result[0], images)
+    assert result[0].result[1]['waveform'].shape[-1] == 66
+    # Exercise real async aggregation without an HTTP endpoint or worker.
+    queue = asyncio.Queue()
+    await queue.put({'worker_id': 'worker_2', 'tensor': images[2:3], 'image_index': 0, 'is_last': True})
+    prompt_server.distributed_pending_jobs['v3-aggregate'] = queue
+    result = await map_node(collector,
+                            {'images': images[:2], 'load_balance': False,
+                             'multi_job_id': 'v3-aggregate', 'enabled_worker_ids': '["worker_2"]'})
+    assert torch.equal(result[0].result[0], images[:3])
+    assert 'v3-aggregate' not in prompt_server.distributed_pending_jobs
+
+    # The registered node must forward the transfer options to the runtime
+    # helper, and default to the legacy transfer when they are not set.
+    collector_runtime = sys.modules[module.__name__ + '.nodes.collector'].DistributedCollectorNode
+    forwarded = []
+    def fake_run(self, **kwargs):
+        forwarded.append(kwargs)
+        return (kwargs['images'], kwargs['audio'])
+    with patch.object(collector_runtime, 'run', fake_run):
+        await map_node(collector, {'images': images[:1], 'load_balance': False, 'pass_through': True,
+                                   'transfer_format': 'av1', 'transfer_quality': 70})
+        await map_node(collector, {'images': images[:1], 'load_balance': False, 'pass_through': True})
+    # Set values arrive list-wrapped (is_input_list); unset ones use the
+    # execute() defaults. The runtime helper unwraps both forms.
+    assert [(kw['transfer_format'], kw['transfer_quality']) for kw in forwarded] == [
+        (['av1'], [70]), ('legacy_png', 90)], forwarded
+
+    # Prove V3's per-execution class clones do not leak helper instance state;
+    # replace only the GPU/model boundary, not the input parser or executor.
+    upscale_runtime = sys.modules[module.__name__ + '.nodes.distributed_upscale'].UltimateSDUpscaleDistributed
+    seen = []
+    def fake_upscale(self, *args):
+        assert not hasattr(self, 'acceptance_marker')
+        self.acceptance_marker = True
+        seen.append(args)
+        return (args[0],)
+    inputs = {'upscaled_image': images, 'model': object(), 'positive': [], 'negative': [],
+              'vae': object(), 'seed': 1, 'steps': 1, 'cfg': 1.0,
+              'sampler_name': 'euler', 'scheduler': 'normal', 'denoise': 0.5,
+              'tile_width': 64, 'tile_height': 64, 'padding': 0, 'mask_blur': 0,
+              'force_uniform_tiles': True, 'tiled_decode': False,
+              'multi_job_id': 'tile-job', 'is_worker': True, 'master_url': 'http://master.invalid',
+              'enabled_worker_ids': '["worker_2"]', 'worker_id': 'worker_2',
+              'tile_indices': '[2,3]', 'dynamic_threshold': 9}
+    with patch.object(upscale_runtime, 'run', fake_upscale):
+        for _ in range(2):
+            result = await map_node(mapping['UltimateSDUpscaleDistributed'], inputs)
+            assert result[0].result[0] is images
+    assert len(seen) == 2 and seen[0][-7:] == tuple(inputs[name] for name in (
+        'multi_job_id', 'is_worker', 'master_url', 'enabled_worker_ids', 'worker_id',
+        'tile_indices', 'dynamic_threshold')), seen[0]
+    import math
+    assert math.isnan(mapping['UltimateSDUpscaleDistributed'].fingerprint_inputs(multi_job_id='tile-job'))
+    assert math.isnan(mapping['UltimateSDUpscaleDistributed'].fingerprint_inputs(multi_job_id=''))
+
+    graph = {
+        '1': {'class_type': 'EmptyImage', 'inputs': {'height': 8, 'width': 8, 'batch_size': 10, 'color': 0}},
+        '2': {'class_type': 'DistributedCollector', 'inputs': {'images': ['1', 0], 'load_balance': False,
+                                                              'pass_through': True}},
+        '3': {'class_type': 'ImageBatchDivider', 'inputs': {'images': ['2', 0], 'divide_by': 10}},
+        '4': {'class_type': 'PreviewImage', 'inputs': {'images': ['3', 9]}},
+    }
+    valid = await execution.validate_prompt('v3-graph', graph, None)
+    assert valid[0], valid
+    invalid = {'1': {'class_type': 'UltimateSDUpscaleDistributed',
+                     'inputs': {**inputs, 'sampler_name': 'INVALID_ENUM', 'scheduler': 'INVALID_ENUM'}},
+               '2': {'class_type': 'PreviewImage', 'inputs': {'images': ['1', 0]}}}
+    rejected = await execution.validate_prompt('v3-invalid', invalid, None)
+    assert not rejected[0]
+    errors = [error for entry in rejected[3].values() for error in entry['errors']]
+    enum_errors = [error['extra_info']['input_name'] for error in errors if error['type'] == 'value_not_in_list']
+    assert {'sampler_name', 'scheduler'} <= set(enum_errors), errors
+    import folder_paths
+    scratch = Path(os.environ.get('TMPDIR', Path.home() / '.hermes/cache/scratch'))
+    with tempfile.TemporaryDirectory(prefix='v3-preview-', dir=scratch) as temp:
+        with patch.object(folder_paths, 'temp_directory', temp):
+            executor = execution.PromptExecutor(
+                prompt_server, cache_args={'ram': 0, 'ram_inactive': 0}, asset_manager=asset_manager)
+            await asyncio.to_thread(executor.execute, graph, 'v3-graph', {}, valid[2])
+            assert executor.success, executor.status_messages
+            history = executor.history_result
+            record = history['outputs']['4']['images'][0]
+            preview = Path(temp) / record.get('subfolder', '') / record['filename']
+            with Image.open(preview) as image:
+                assert image.size == (8, 8) and image.mode == 'RGB'
+                assert image.getextrema() == ((0, 0), (0, 0), (0, 0))
+
+    # A worker prompt terminated by the sink validates and runs through the
+    # real executor without writing anything to the temp directory.
+    sink_graph = {
+        '1': {'class_type': 'EmptyImage', 'inputs': {'height': 8, 'width': 8, 'batch_size': 4, 'color': 0}},
+        '2': {'class_type': 'DistributedCollector', 'inputs': {'images': ['1', 0], 'load_balance': False,
+                                                              'pass_through': True}},
+        '3': {'class_type': 'DistributedWorkerSink', 'inputs': {'value': ['2', 0]}},
+    }
+    valid = await execution.validate_prompt('v3-sink', sink_graph, None)
+    assert valid[0], valid
+    assert valid[2] == ['3'], valid[2]
+    collected = []
+    def record_run(self, **kwargs):
+        collected.append(kwargs['images'])
+        return (kwargs['images'][0], None)
+    with tempfile.TemporaryDirectory(prefix='v3-sink-', dir=scratch) as temp:
+        with patch.object(folder_paths, 'temp_directory', temp), patch.object(collector_runtime, 'run', record_run):
+            executor = execution.PromptExecutor(
+                prompt_server, cache_args={'ram': 0, 'ram_inactive': 0}, asset_manager=asset_manager)
+            await asyncio.to_thread(executor.execute, sink_graph, 'v3-sink', {}, valid[2])
+            assert executor.success, executor.status_messages
+            assert len(collected) == 1 and collected[0][0].shape == (4, 8, 8, 3), collected
+            assert list(Path(temp).rglob('*')) == [], list(Path(temp).rglob('*'))
+            assert executor.history_result['outputs'].get('3', {}) == {}, executor.history_result
+    print('V3_EXECUTION_OK nine nodes; upscale GPU boundary mocked; preview artifact verified; sink writes nothing')
+
+
+async def main():
+    asset_manager = default_asset_manager()
+    prompt_server = server.PromptServer(asyncio.get_running_loop(), asset_manager)
+    assert not (ROOT / 'distributed.py').exists(), 'obsolete root bootstrap remains'
+    assert await nodes.load_custom_node(str(ROOT)), 'ComfyUI loader rejected the pack'
+    module = sys.modules[str(ROOT).replace('.', '_x_')]
+    assert not hasattr(module, 'NODE_CLASS_MAPPINGS'), 'V1 map shadows V3 entrypoint'
+    extension = await module.comfy_entrypoint()
+    classes = await extension.get_node_list()
+    mapping = {cls.GET_SCHEMA().node_id: cls for cls in classes}
+    assert len(classes) == len(mapping) == len(BASELINE) + len(ADDED_NODES) == 9
+    assert set(mapping) == set(BASELINE) | ADDED_NODES
+    for node_id in ADDED_NODES:
+        assert nodes.NODE_CLASS_MAPPINGS[node_id] is mapping[node_id]
+    sink_info = mapping['DistributedWorkerSink'].GET_NODE_INFO_V1()
+    assert sink_info['output_node'] and sink_info['dev_only'], sink_info
+    for node_id, cls in mapping.items():
+        if node_id in ADDED_NODES:
+            continue
+        assert issubclass(cls, io.ComfyNode)
+        assert nodes.NODE_CLASS_MAPPINGS[node_id] is cls
+        old = dict(BASELINE[node_id])
+        if node_id in ('ImageBatchDivider', 'AudioBatchDivider'):
+            # V1's ByPassTypeTuple advertises '*' when indexed, while its
+            # underlying tuple and existing frontend declare IMAGE/AUDIO.
+            # Native V3 declares all ten existing typed sockets explicitly.
+            assert old['output'] == ['*'] * 10
+            old['output'] = ['IMAGE' if node_id == 'ImageBatchDivider' else 'AUDIO'] * 10
+        new = json.loads(json.dumps(cls.GET_NODE_INFO_V1()))
+        for group in ('required', 'optional'):
+            old_inputs = old['input'].get(group, {})
+            new_inputs = dict(new['input'].get(group, {}))
+            if group == 'optional':
+                for name, expected in ADDED_OPTIONAL_INPUTS.get(node_id, {}).items():
+                    assert name in new_inputs, (node_id, name, 'added input missing')
+                    actual = normalized_input(new_inputs.pop(name))
+                    kind, opts = normalized_input(expected)
+                    assert actual[0] == kind, (node_id, name, actual)
+                    for key, value in opts.items():
+                        assert actual[1].get(key) == value, (node_id, name, key, actual)
+            assert list(old_inputs) == list(new_inputs), (node_id, group, 'input order')
+            for name, original in old_inputs.items():
+                assert normalized_input(original) == normalized_input(new_inputs[name]), (node_id, name, original, new_inputs[name])
+        for key in ('output', 'output_name', 'output_is_list', 'is_input_list', 'output_node', 'category', 'display_name'):
+            assert old[key] == new[key], (node_id, key, old[key], new[key])
+        # Standard context lives in cls.hidden; orchestrator metadata remains
+        # accepted by its original kwarg name, without creating new widgets.
+        signature = inspect.signature(cls.execute)
+        for name, field in old['input'].get('hidden', {}).items():
+            if isinstance(field, list):
+                assert cls.GET_SCHEMA().accept_all_inputs, node_id
+                assert name in signature.parameters, (node_id, name)
+                assert signature.parameters[name].default == field[1]['default'], (node_id, name)
+            else:
+                assert name in new['input']['hidden'], (node_id, name)
+        expected_hidden = []
+        if node_id == 'DistributedModelName':
+            expected_hidden.extend(['unique_id', 'extra_pnginfo'])
+        if old['output_node']:
+            expected_hidden.extend(name for name in ['prompt', 'extra_pnginfo'] if name not in expected_hidden)
+        assert list(new['input'].get('hidden', {})) == expected_hidden, (node_id, new['input'].get('hidden'))
+    print('SCHEMA_PARITY_OK', len(mapping))
+    check_saved_workflows(mapping)
+    await check_execution(mapping, module, prompt_server, asset_manager)
+    print('V3_ACCEPTANCE_OK')
+
+
+asyncio.run(main())

tests/test_async_helpers.py

@@ -1,4 +1,3 @@
-import asyncio
 import importlib.util
 import sys
 import types
@@ -6,7 +5,7 @@ import unittest
 from pathlib import Path
 
 
-class _FakePromptQueue:
+class _PromptQueue:
     def __init__(self):
         self.items = []
 
@@ -14,16 +13,10 @@ class _FakePromptQueue:
         self.items.append(item)
 
 
-class _FakePromptServer:
-    def __init__(self):
-        self.number = 0
-        self.prompt_queue = _FakePromptQueue()
+def _load_async_helpers_module():
+    module_path = Path(__file__).resolve().parents[1] / "utils" / "async_helpers.py"
+    package_name = "dist_async_helpers_testpkg"
 
-    def trigger_on_prompt(self, payload):
-        return payload
-
-
-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]
@@ -36,68 +29,62 @@ def _bootstrap_package(package_name):
     utils_pkg.__path__ = []
     sys.modules[f"{package_name}.utils"] = utils_pkg
 
-    network_module = types.ModuleType(f"{package_name}.utils.network")
-    network_module.get_server_loop = asyncio.get_event_loop
-    sys.modules[f"{package_name}.utils.network"] = network_module
-
-    fake_prompt_server = _FakePromptServer()
-
-    server_module = types.ModuleType("server")
-    server_module.PromptServer = types.SimpleNamespace(instance=fake_prompt_server)
-    sys.modules["server"] = server_module
-
     execution_module = types.ModuleType("execution")
 
-    async def _validate_prompt(_prompt_id, _prompt, _partial_targets):
-        return (True, None, ["1"], {})
+    async def _validate_prompt(prompt_id, prompt, partial_execution_targets):
+        return (True, None, ["9"], {})
 
     execution_module.validate_prompt = _validate_prompt
-    execution_module.SENSITIVE_EXTRA_DATA_KEYS = ()
+    execution_module.SENSITIVE_EXTRA_DATA_KEYS = []
     sys.modules["execution"] = execution_module
 
-    return fake_prompt_server
-
-
-def _load_async_helpers_module():
-    package_name = "dist_async_helpers_testpkg"
-    fake_prompt_server = _bootstrap_package(package_name)
-    module_path = Path(__file__).resolve().parents[3] / "utils/async_helpers.py"
-    spec = importlib.util.spec_from_file_location(
-        f"{package_name}.utils.async_helpers",
-        module_path,
+    prompt_server = types.SimpleNamespace(
+        trigger_on_prompt=lambda payload: payload,
+        number=12,
+        prompt_queue=_PromptQueue(),
     )
+    server_module = types.ModuleType("server")
+    server_module.PromptServer = types.SimpleNamespace(instance=prompt_server)
+    sys.modules["server"] = server_module
+
+    network_module = types.ModuleType(f"{package_name}.utils.network")
+    network_module.get_server_loop = lambda: None
+    sys.modules[f"{package_name}.utils.network"] = network_module
+
+    spec = importlib.util.spec_from_file_location(f"{package_name}.utils.async_helpers", 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, fake_prompt_server
+    return module, prompt_server
 
 
-class AsyncHelpersQueuePromptPayloadTests(unittest.TestCase):
-    def test_queue_prompt_payload_includes_create_time_metadata(self):
-        async_helpers_module, fake_prompt_server = _load_async_helpers_module()
+async_helpers, prompt_server = _load_async_helpers_module()
 
-        prompt_id = asyncio.run(
-            async_helpers_module.queue_prompt_payload(
-                {"1": {"class_type": "KSampler", "inputs": {}}},
-                workflow_meta={"id": "workflow-1"},
-                client_id="client-1",
-            )
+
+class QueuePromptPayloadTests(unittest.IsolatedAsyncioTestCase):
+    async def test_queue_prompt_payload_includes_create_time_and_client_metadata(self):
+        result = await async_helpers.queue_prompt_payload(
+            {"1": {"class_type": "Node"}},
+            workflow_meta={"id": "workflow-1"},
+            client_id="client-1",
+            include_queue_metadata=True,
         )
 
-        self.assertIsInstance(prompt_id, str)
-        self.assertTrue(fake_prompt_server.prompt_queue.items)
-
-        queued_item = fake_prompt_server.prompt_queue.items[-1]
-        self.assertEqual(len(queued_item), 6)
-        self.assertEqual(queued_item[1], prompt_id)
+        self.assertIsInstance(result["prompt_id"], str)
+        self.assertTrue(result["prompt_id"])
+        self.assertEqual(result["number"], 12)
+        self.assertEqual(result["node_errors"], {})
 
+        self.assertEqual(prompt_server.number, 13)
+        self.assertEqual(len(prompt_server.prompt_queue.items), 1)
+        queued_item = prompt_server.prompt_queue.items[0]
+        self.assertEqual(queued_item[0], 12)
         extra_data = queued_item[3]
         self.assertEqual(extra_data["client_id"], "client-1")
-        self.assertEqual(extra_data["extra_pnginfo"]["workflow"]["id"], "workflow-1")
         self.assertIn("create_time", extra_data)
         self.assertIsInstance(extra_data["create_time"], int)
         self.assertGreater(extra_data["create_time"], 0)
+        self.assertEqual(extra_data["extra_pnginfo"]["workflow"], {"id": "workflow-1"})
 
 
 if __name__ == "__main__":

tests/test_batch_dividers.py

@@ -8,7 +8,7 @@ import torch
 
 
 def _load_utilities_module():
-    module_path = Path(__file__).resolve().parents[3] / "nodes" / "utilities.py"
+    module_path = Path(__file__).resolve().parents[1] / "nodes" / "utilities.py"
     package_name = "dist_divider_testpkg"
 
     for mod_name in list(sys.modules):

tests/test_collector_list_inputs.py

@@ -0,0 +1,630 @@
+import importlib.util
+import sys
+import types
+import asyncio
+from pathlib import Path
+
+import torch
+
+
+def _register_real_transfer_codec(package_name, utils_pkg):
+    """Load the real utils/transfer_codec.py into the synthetic test package."""
+    codec_path = Path(__file__).resolve().parents[1] / "utils" / "transfer_codec.py"
+    codec_spec = importlib.util.spec_from_file_location(f"{package_name}.utils.transfer_codec", codec_path)
+    codec_module = importlib.util.module_from_spec(codec_spec)
+    sys.modules[f"{package_name}.utils.transfer_codec"] = codec_module
+    codec_spec.loader.exec_module(codec_module)
+    utils_pkg.transfer_codec = codec_module
+    return codec_module
+
+
+def _load_collector_module():
+    module_path = Path(__file__).resolve().parents[1] / "nodes" / "collector.py"
+    package_name = "dist_collector_list_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
+
+    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
+    _register_real_transfer_codec(package_name, utils_pkg)
+
+    class _Routes:
+        def post(self, _path):
+            return lambda fn: fn
+
+        def get(self, _path):
+            return lambda fn: fn
+
+    prompt_server = types.SimpleNamespace(
+        routes=_Routes(),
+        distributed_jobs_lock=None,
+        distributed_pending_jobs={},
+    )
+    server_module = types.ModuleType("server")
+    server_module.PromptServer = types.SimpleNamespace(instance=prompt_server)
+    sys.modules["server"] = server_module
+
+    comfy_module = types.ModuleType("comfy")
+    model_management = types.ModuleType("comfy.model_management")
+
+    class InterruptProcessingException(Exception):
+        pass
+
+    model_management.InterruptProcessingException = InterruptProcessingException
+    model_management.throw_exception_if_processing_interrupted = lambda: None
+    comfy_module.model_management = model_management
+
+    comfy_utils = types.ModuleType("comfy.utils")
+
+    class ProgressBar:
+        def __init__(self, _total):
+            self.total = _total
+            self.updates = []
+
+        def update(self, value):
+            self.updates.append(value)
+
+    comfy_utils.ProgressBar = ProgressBar
+    comfy_module.utils = comfy_utils
+    sys.modules["comfy"] = comfy_module
+    sys.modules["comfy.model_management"] = model_management
+    sys.modules["comfy.utils"] = comfy_utils
+
+    aiohttp_module = types.ModuleType("aiohttp")
+    aiohttp_module.ClientTimeout = lambda total: types.SimpleNamespace(total=total)
+    sys.modules["aiohttp"] = aiohttp_module
+
+    logging_module = types.ModuleType(f"{package_name}.utils.logging")
+    logging_module.debug_log = lambda *_args, **_kwargs: None
+    logging_module.log = lambda *_args, **_kwargs: None
+    sys.modules[f"{package_name}.utils.logging"] = logging_module
+
+    config_module = types.ModuleType(f"{package_name}.utils.config")
+    config_module.get_worker_timeout_seconds = lambda: 0.1
+    config_module.load_config = lambda: {"workers": []}
+    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
+    constants_module.TRANSFER_MAX_CHUNK_BYTES = 32 * 1024 * 1024
+    constants_module.TRANSFER_MAX_REQUEST_BYTES = 96 * 1024 * 1024
+    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() if hasattr(tensor, "contiguous") else tensor
+
+    image_module.ensure_contiguous = _ensure_contiguous
+    image_module.tensor_to_pil = lambda *_args, **_kwargs: None
+    image_module.pil_to_tensor = lambda value: value
+    sys.modules[f"{package_name}.utils.image"] = image_module
+
+    network_module = types.ModuleType(f"{package_name}.utils.network")
+    network_module.build_worker_url = lambda worker: "http://worker"
+    network_module.get_client_session = lambda: None
+    network_module.probe_worker = lambda *_args, **_kwargs: None
+    sys.modules[f"{package_name}.utils.network"] = network_module
+
+    audio_payload_module = types.ModuleType(f"{package_name}.utils.audio_payload")
+    audio_payload_module.encode_audio_payload = lambda audio: audio
+    sys.modules[f"{package_name}.utils.audio_payload"] = audio_payload_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
+
+    spec = importlib.util.spec_from_file_location(f"{package_name}.nodes.collector", module_path)
+    module = importlib.util.module_from_spec(spec)
+    assert spec is not None and spec.loader is not None
+    spec.loader.exec_module(module)
+    return module
+
+
+def test_collector_opts_into_comfyui_list_inputs():
+    collector = _load_collector_module().DistributedCollectorNode
+
+    assert collector.INPUT_IS_LIST is True
+
+
+def test_collector_exposes_images_as_optional_input():
+    input_types = _load_collector_module().DistributedCollectorNode.INPUT_TYPES()
+
+    assert "images" not in input_types["required"]
+    assert input_types["optional"]["images"] == ("IMAGE",)
+
+
+def test_audio_only_pass_through_returns_no_images_and_preserves_audio():
+    collector = _load_collector_module().DistributedCollectorNode()
+    audio = {"waveform": torch.ones(1, 2, 4), "sample_rate": 48000}
+
+    images, returned_audio = collector.run(images=None, audio=[audio])
+
+    assert images is None
+    assert returned_audio is audio
+
+
+def test_collector_rejects_missing_images_and_audio():
+    collector = _load_collector_module().DistributedCollectorNode()
+
+    try:
+        collector.run(images=None, audio=None)
+    except ValueError as exc:
+        assert "image or audio" in str(exc).lower()
+    else:
+        raise AssertionError("Expected collector to reject a run with no media input")
+
+
+def test_delegate_only_master_allows_no_local_media_input():
+    collector = _load_collector_module().DistributedCollectorNode()
+
+    images, audio = collector.run(
+        images=None,
+        audio=None,
+        multi_job_id=["delegate-audio-job"],
+        delegate_only=[True],
+        enabled_worker_ids=["[]"],
+    )
+
+    assert images is None
+    assert tuple(audio["waveform"].shape) == (1, 2, 1)
+
+
+def test_pass_through_collapses_comfyui_image_list_to_batch_and_unwraps_hidden_inputs():
+    collector = _load_collector_module().DistributedCollectorNode()
+    first = torch.zeros(1, 2, 2, 3)
+    second = torch.ones(1, 2, 2, 3)
+
+    images, audio = collector.run(
+        images=[first, second],
+        load_balance=[False],
+        audio=[None],
+        multi_job_id=[""],
+        is_worker=[False],
+        master_url=[""],
+        enabled_worker_ids=["[]"],
+        worker_batch_size=[1],
+        worker_id=[""],
+        pass_through=[False],
+        delegate_only=[False],
+    )
+
+    assert tuple(images.shape) == (2, 2, 2, 3)
+    assert torch.equal(images[0:1], first)
+    assert torch.equal(images[1:2], second)
+    assert tuple(audio["waveform"].shape) == (1, 2, 1)
+
+
+def test_worker_list_input_sends_one_completion_sequence_with_last_only_on_final_item():
+    module = _load_collector_module()
+    collector = module.DistributedCollectorNode()
+    first = torch.zeros(1, 2, 2, 3)
+    second = torch.ones(1, 2, 2, 3)
+    posted_payloads = []
+
+    class _FakeImage:
+        def save(self, fp, format=None, compress_level=None):
+            fp.write(b"png-bytes")
+
+    class _FakeResponse:
+        async def __aenter__(self):
+            return self
+
+        async def __aexit__(self, exc_type, exc, tb):
+            return False
+
+        def raise_for_status(self):
+            return None
+
+    class _FakeSession:
+        def post(self, url, json, timeout):
+            posted_payloads.append(json)
+            return _FakeResponse()
+
+    async def _fake_get_client_session():
+        return _FakeSession()
+
+    module.tensor_to_pil = lambda *_args, **_kwargs: _FakeImage()
+    module.get_client_session = _fake_get_client_session
+    module.encode_audio_payload = lambda _audio: None
+
+    images, audio = collector.run(
+        images=[first, second],
+        load_balance=[False],
+        audio=[None],
+        multi_job_id=["job-list-1"],
+        is_worker=[True],
+        master_url=["http://master"],
+        enabled_worker_ids=["[]"],
+        worker_batch_size=[1],
+        worker_id=["worker-a"],
+        pass_through=[False],
+        delegate_only=[False],
+    )
+
+    assert tuple(images.shape) == (2, 2, 2, 3)
+    assert tuple(audio["waveform"].shape) == (1, 2, 1)
+    assert len(posted_payloads) == 2
+    assert [payload["batch_idx"] for payload in posted_payloads] == [0, 1]
+    assert [payload["is_last"] for payload in posted_payloads] == [False, True]
+    assert {payload["job_id"] for payload in posted_payloads} == {"job-list-1"}
+    assert {payload["worker_id"] for payload in posted_payloads} == {"worker-a"}
+
+
+def test_audio_only_worker_sends_one_completion_without_image():
+    module = _load_collector_module()
+    collector = module.DistributedCollectorNode()
+    audio = {"waveform": torch.ones(1, 2, 4), "sample_rate": 48000}
+    posted = []
+
+    class _FakeResponse:
+        async def __aenter__(self):
+            return self
+
+        async def __aexit__(self, exc_type, exc, tb):
+            return False
+
+        def raise_for_status(self):
+            return None
+
+    class _FakeSession:
+        def post(self, url, json, timeout):
+            posted.append((url, json, timeout.total))
+            return _FakeResponse()
+
+    async def _fake_get_client_session():
+        return _FakeSession()
+
+    module.get_client_session = _fake_get_client_session
+    module.encode_audio_payload = lambda value: {"encoded": value is audio}
+
+    images, returned_audio = collector.run(
+        images=None,
+        audio=[audio],
+        multi_job_id=["audio-job"],
+        is_worker=[True],
+        master_url=["http://master"],
+        worker_id=["worker-a"],
+    )
+
+    assert images is None
+    assert returned_audio is audio
+    assert len(posted) == 1
+    assert posted[0][0] == "http://master/distributed/job_complete"
+    assert posted[0][2] == 600
+    assert posted[0][1] == {
+        "job_id": "audio-job",
+        "worker_id": "worker-a",
+        "batch_idx": 0,
+        "audio": {"encoded": True},
+        "is_last": True,
+    }
+
+
+def test_audio_only_master_combines_local_and_worker_audio():
+    module = _load_collector_module()
+    collector = module.DistributedCollectorNode()
+    master_audio = {"waveform": torch.ones(1, 2, 2), "sample_rate": 48000}
+    worker_audio = {"waveform": torch.full((1, 2, 3), 2.0), "sample_rate": 48000}
+    module.prompt_server.distributed_jobs_lock = asyncio.Lock()
+    queue = asyncio.Queue()
+    queue.put_nowait(
+        {
+            "worker_id": "worker-a",
+            "image_index": 0,
+            "tensor": None,
+            "audio": worker_audio,
+            "is_last": True,
+        }
+    )
+    module.prompt_server.distributed_pending_jobs = {"audio-job": queue}
+
+    images, combined_audio = asyncio.run(
+        collector.execute(
+            images=None,
+            audio=master_audio,
+            multi_job_id="audio-job",
+            enabled_worker_ids='["worker-a"]',
+        )
+    )
+
+    assert images is None
+    assert combined_audio["sample_rate"] == 48000
+    assert tuple(combined_audio["waveform"].shape) == (1, 2, 5)
+    assert torch.equal(combined_audio["waveform"][..., :2], master_audio["waveform"])
+    assert torch.equal(combined_audio["waveform"][..., 2:], worker_audio["waveform"])
+
+
+def test_delegate_only_audio_collects_worker_audio_without_placeholder_image():
+    module = _load_collector_module()
+    collector = module.DistributedCollectorNode()
+    worker_audio = {"waveform": torch.full((1, 2, 3), 2.0), "sample_rate": 48000}
+    module.prompt_server.distributed_jobs_lock = asyncio.Lock()
+    queue = asyncio.Queue()
+    queue.put_nowait(
+        {
+            "worker_id": "worker-a",
+            "image_index": 0,
+            "tensor": None,
+            "audio": worker_audio,
+            "is_last": True,
+        }
+    )
+    module.prompt_server.distributed_pending_jobs = {"delegate-audio-job": queue}
+
+    images, combined_audio = asyncio.run(
+        collector.execute(
+            images=None,
+            audio=None,
+            multi_job_id="delegate-audio-job",
+            enabled_worker_ids='["worker-a"]',
+            delegate_only=True,
+        )
+    )
+
+    assert images is None
+    assert combined_audio["sample_rate"] == 48000
+    assert torch.equal(combined_audio["waveform"], worker_audio["waveform"])
+
+
+# --------------------------------------------------------------------------- #
+# Compressed worker -> master transfer
+# --------------------------------------------------------------------------- #
+
+class _RecordingResponse:
+    def __init__(self, status=200, text=""):
+        self.status = status
+        self._text = text
+
+    async def __aenter__(self):
+        return self
+
+    async def __aexit__(self, exc_type, exc, tb):
+        return False
+
+    async def text(self):
+        return self._text
+
+    def raise_for_status(self):
+        if self.status >= 400:
+            raise RuntimeError(f"HTTP {self.status}")
+
+
+class _RecordingSession:
+    def __init__(self, statuses=None, text="Invalid JSON payload"):
+        self.calls = []
+        self._statuses = list(statuses or [])
+        self._text = text
+
+    def post(self, url, timeout, json=None, data=None):
+        self.calls.append({"url": url, "json": json, "data": data, "timeout": timeout.total})
+        status = self._statuses.pop(0) if self._statuses else 200
+        return _RecordingResponse(status=status, text=self._text)
+
+
+def _install_fake_transport(module, session):
+    async def _fake_get_client_session():
+        return session
+
+    module.get_client_session = _fake_get_client_session
+    # Record the multipart content instead of building a real aiohttp.FormData.
+    module.DistributedCollectorNode._build_multipart = staticmethod(
+        lambda envelope, parts: {"envelope": envelope, "parts": list(parts)}
+    )
+
+
+def _run_worker(collector, images, audio=None, **kwargs):
+    return collector.run(
+        images=[images],
+        audio=[audio],
+        multi_job_id=["job-c"],
+        is_worker=[True],
+        master_url=["http://master"],
+        worker_id=["worker-a"],
+        **{key: [value] for key, value in kwargs.items()},
+    )
+
+
+def test_collector_exposes_transfer_widgets_with_legacy_default():
+    module = _load_collector_module()
+    optional = module.DistributedCollectorNode.INPUT_TYPES()["optional"]
+
+    formats, options = optional["transfer_format"]
+    assert "legacy_png" in formats and "h264" in formats and "ffv1" in formats
+    assert options["default"] == "legacy_png"
+    assert optional["transfer_quality"][0] == "INT"
+
+
+def test_worker_sends_video_batch_as_one_multipart_chunk_with_audio_on_last():
+    module = _load_collector_module()
+    collector = module.DistributedCollectorNode()
+    session = _RecordingSession()
+    _install_fake_transport(module, session)
+    module.encode_audio_payload = lambda value: None if value is None else {"encoded": True}
+    images = torch.rand(6, 32, 48, 3)
+
+    _run_worker(
+        collector,
+        images,
+        audio={"waveform": torch.ones(1, 2, 4), "sample_rate": 48000},
+        transfer_format="ffv1",
+        transfer_quality=90,
+    )
+
+    assert len(session.calls) == 1
+    call = session.calls[0]
+    assert call["url"] == "http://master/distributed/job_complete"
+    assert call["json"] is None
+    envelope = call["data"]["envelope"]
+    assert envelope["format"] == "ffv1"
+    assert (envelope["start_index"], envelope["count"]) == (0, 6)
+    assert (envelope["width"], envelope["height"]) == (48, 32)
+    assert envelope["is_last"] is True
+    assert envelope["audio"] == {"encoded": True}
+    assert envelope["job_id"] == "job-c" and envelope["worker_id"] == "worker-a"
+    assert len(call["data"]["parts"]) == 1
+
+
+def test_worker_marks_only_final_chunk_as_last():
+    module = _load_collector_module()
+    collector = module.DistributedCollectorNode()
+    session = _RecordingSession()
+    _install_fake_transport(module, session)
+    module.TRANSFER_MAX_CHUNK_BYTES = 1  # force one frame per request
+
+    _run_worker(collector, torch.rand(3, 8, 8, 3), transfer_format="png")
+
+    envelopes = [call["data"]["envelope"] for call in session.calls]
+    assert [(e["start_index"], e["count"], e["is_last"]) for e in envelopes] == [
+        (0, 1, False),
+        (1, 1, False),
+        (2, 1, True),
+    ]
+    assert "audio" not in envelopes[-1]
+
+
+def test_worker_falls_back_to_legacy_png_when_master_rejects_multipart():
+    from PIL import Image
+    import numpy as np
+
+    module = _load_collector_module()
+    collector = module.DistributedCollectorNode()
+    session = _RecordingSession(statuses=[400])
+    _install_fake_transport(module, session)
+    module.tensor_to_pil = lambda tensor, idx: Image.fromarray(
+        (255 * tensor[idx].cpu().numpy()).astype(np.uint8)
+    )
+
+    _run_worker(collector, torch.rand(2, 8, 8, 3), transfer_format="h264")
+
+    assert session.calls[0]["data"] is not None  # compressed attempt
+    legacy = [call["json"] for call in session.calls[1:]]
+    assert [payload["batch_idx"] for payload in legacy] == [0, 1]
+    assert all(payload["image"].startswith("data:image/png;base64,") for payload in legacy)
+    assert [payload["is_last"] for payload in legacy] == [False, True]
+
+
+def test_worker_uses_legacy_path_for_unknown_transfer_format():
+    from PIL import Image
+    import numpy as np
+
+    module = _load_collector_module()
+    collector = module.DistributedCollectorNode()
+    session = _RecordingSession()
+    _install_fake_transport(module, session)
+    module.tensor_to_pil = lambda tensor, idx: Image.fromarray(
+        (255 * tensor[idx].cpu().numpy()).astype(np.uint8)
+    )
+
+    _run_worker(collector, torch.rand(1, 8, 8, 3), transfer_format="gif")
+
+    assert len(session.calls) == 1
+    assert session.calls[0]["json"]["image"].startswith("data:image/png;base64,")
+
+
+def test_worker_falls_back_when_video_encoder_missing():
+    module = _load_collector_module()
+    collector = module.DistributedCollectorNode()
+    session = _RecordingSession()
+    _install_fake_transport(module, session)
+    codec = module.transfer_codec
+    original = codec.video_codec_available
+    codec.video_codec_available = lambda fmt: fmt == "h264" and original("h264")
+    try:
+        _run_worker(collector, torch.rand(2, 16, 16, 3), transfer_format="av1")
+    finally:
+        codec.video_codec_available = original
+
+    assert session.calls[0]["data"]["envelope"]["format"] == "h264"
+
+
+def test_worker_does_not_fall_back_on_non_legacy_rejection():
+    import pytest
+
+    module = _load_collector_module()
+    collector = module.DistributedCollectorNode()
+    session = _RecordingSession(statuses=[400], text="Invalid multipart payload: Request Entity Too Large")
+    _install_fake_transport(module, session)
+
+    with pytest.raises(RuntimeError, match="MAX_REQUEST_BYTES"):
+        _run_worker(collector, torch.rand(2, 8, 8, 3), transfer_format="png")
+
+    # One compressed attempt, no legacy resend of the batch.
+    assert len(session.calls) == 1
+    assert session.calls[0]["json"] is None
+
+
+def test_worker_sends_nothing_when_a_frame_cannot_fit():
+    import numpy as np
+    import pytest
+
+    module = _load_collector_module()
+    collector = module.DistributedCollectorNode()
+    session = _RecordingSession()
+    _install_fake_transport(module, session)
+    reserve = module.DistributedCollectorNode.TRANSFER_ENVELOPE_RESERVE_BYTES
+    module.TRANSFER_MAX_REQUEST_BYTES = reserve + 32 * 1024
+    module.TRANSFER_MAX_CHUNK_BYTES = 32 * 1024
+    # Frame 0 is tiny; frame 2 is noise that cannot fit in 32 KiB as PNG.
+    images = torch.zeros(3, 128, 128, 3)
+    images[2] = torch.from_numpy(np.random.default_rng(1).random((128, 128, 3), dtype=np.float32))
+
+    with pytest.raises(module.transfer_codec.FrameTooLargeError):
+        _run_worker(collector, images, transfer_format="png")
+
+    assert session.calls == []
+
+
+def test_worker_sends_audio_separately_when_it_would_overflow_last_request():
+    module = _load_collector_module()
+    collector = module.DistributedCollectorNode()
+    session = _RecordingSession()
+    _install_fake_transport(module, session)
+    reserve = module.DistributedCollectorNode.TRANSFER_ENVELOPE_RESERVE_BYTES
+    module.TRANSFER_MAX_REQUEST_BYTES = reserve + 64 * 1024
+    module.TRANSFER_MAX_CHUNK_BYTES = 64 * 1024
+    module.encode_audio_payload = lambda value: None if value is None else {"data": "x" * (128 * 1024)}
+
+    _run_worker(
+        collector,
+        torch.zeros(2, 16, 16, 3),
+        audio={"waveform": torch.ones(1, 2, 4), "sample_rate": 48000},
+        transfer_format="png",
+    )
+
+    multipart = [call["data"]["envelope"] for call in session.calls if call["data"] is not None]
+    assert [e["is_last"] for e in multipart] == [False]
+    assert all("audio" not in e for e in multipart)
+    final = session.calls[-1]["json"]
+    assert final["is_last"] is True and final["audio"]["data"].startswith("x")
+    assert "image" not in final
+
+
+def test_worker_attaches_audio_to_last_chunk_when_it_fits():
+    module = _load_collector_module()
+    collector = module.DistributedCollectorNode()
+    session = _RecordingSession()
+    _install_fake_transport(module, session)
+    module.encode_audio_payload = lambda value: None if value is None else {"data": "small"}
+
+    _run_worker(
+        collector,
+        torch.zeros(2, 16, 16, 3),
+        audio={"waveform": torch.ones(1, 2, 4), "sample_rate": 48000},
+        transfer_format="png",
+    )
+
+    assert len(session.calls) == 1
+    envelope = session.calls[0]["data"]["envelope"]
+    assert envelope["is_last"] is True and envelope["audio"] == {"data": "small"}

tests/test_config.py

@@ -10,7 +10,7 @@ from unittest.mock import patch
 
 
 def _load_config_module():
-    module_path = Path(__file__).resolve().parents[3] / "utils" / "config.py"
+    module_path = Path(__file__).resolve().parents[1] / "utils" / "config.py"
     package_name = "dist_cfg_testpkg"
 
     for mod_name in list(sys.modules):
@@ -18,16 +18,16 @@ def _load_config_module():
             del sys.modules[mod_name]
 
     root_pkg = types.ModuleType(package_name)
-    root_pkg.__path__ = []
+    root_pkg.__path__ = [str(module_path.parent)]
     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)
@@ -213,7 +213,7 @@ class SaveConfigTests(unittest.TestCase):
                 data = config.get_default_config()
                 config.save_config(data)
                 # Cache is now None; load_config should re-read
-                self.assertIsNone(config._config_state().cache)
+                self.assertIsNone(config._config_cache)
 
     def test_written_file_is_valid_json(self):
         data = config.get_default_config()

tests/test_detection.py

@@ -7,7 +7,7 @@ from unittest.mock import patch
 
 
 def _load_detection_module():
-    module_path = Path(__file__).resolve().parents[3] / "workers" / "detection.py"
+    module_path = Path(__file__).resolve().parents[1] / "workers" / "detection.py"
     package_name = "dist_det_testpkg"
 
     for mod_name in list(sys.modules):
@@ -32,7 +32,6 @@ def _load_detection_module():
 
     network_module = types.ModuleType(f"{package_name}.utils.network")
     network_module.normalize_host = lambda value: value
-    network_module.build_worker_url = lambda worker, endpoint="": f"http://{worker.get('host', 'localhost')}{endpoint}"
 
     async def _fake_session():
         raise RuntimeError("network calls not used in these tests")
@@ -157,7 +156,7 @@ class IsLocalWorkerTests(unittest.IsolatedAsyncioTestCase):
         self.assertTrue(result)
 
     async def test_true_for_0_0_0_0(self):
-        result = await detection.is_local_worker({"host": "0.0.0.0", "port": 8188})  # nosec B104 - explicit wildcard host test case
+        result = await detection.is_local_worker({"host": "0.0.0.0", "port": 8188})
         self.assertTrue(result)
 
     async def test_true_when_type_is_local(self):

tests/test_dispatch_selection.py

@@ -8,7 +8,7 @@ from unittest.mock import patch
 
 
 def _load_dispatch_module():
-    module_path = Path(__file__).resolve().parents[3] / "api" / "orchestration" / "dispatch.py"
+    module_path = Path(__file__).resolve().parents[1] / "api" / "orchestration" / "dispatch.py"
 
     package_name = "dist_dispatch_testpkg"
     root_pkg = types.ModuleType(package_name)
@@ -249,50 +249,6 @@ class DispatchSelectionTests(unittest.IsolatedAsyncioTestCase):
 
         self.assertIsNone(selected)
 
-    async def test_rank_workers_by_load_returns_sorted_by_queue_depth(self):
-        workers = [
-            {"id": "w1", "name": "Worker 1"},
-            {"id": "w2", "name": "Worker 2"},
-            {"id": "w3", "name": "Worker 3"},
-        ]
-        queue_map = {"w1": 4, "w2": 0, "w3": 2}
-
-        async def fake_probe(worker_url, timeout=3.0):
-            worker_id = worker_url.rsplit("/", 1)[-1]
-            return {"exec_info": {"queue_remaining": queue_map[worker_id]}}
-
-        with patch.object(dispatch, "build_worker_url", side_effect=lambda worker: f"http://host/{worker['id']}"), patch.object(
-            dispatch,
-            "probe_worker",
-            side_effect=fake_probe,
-        ):
-            ranked = await dispatch.rank_workers_by_load(workers, probe_concurrency=3)
-
-        self.assertEqual([worker["id"] for worker in ranked], ["w2", "w3", "w1"])
-
-    async def test_rank_workers_by_load_handles_unreachable_workers(self):
-        workers = [
-            {"id": "w1", "name": "Worker 1"},
-            {"id": "w2", "name": "Worker 2"},
-            {"id": "w3", "name": "Worker 3"},
-        ]
-
-        async def fake_probe(worker_url, timeout=3.0):
-            worker_id = worker_url.rsplit("/", 1)[-1]
-            if worker_id == "w2":
-                return None
-            queue_map = {"w1": 1, "w3": 0}
-            return {"exec_info": {"queue_remaining": queue_map[worker_id]}}
-
-        with patch.object(dispatch, "build_worker_url", side_effect=lambda worker: f"http://host/{worker['id']}"), patch.object(
-            dispatch,
-            "probe_worker",
-            side_effect=fake_probe,
-        ):
-            ranked = await dispatch.rank_workers_by_load(workers, probe_concurrency=3)
-
-        self.assertEqual([worker["id"] for worker in ranked], ["w3", "w1", "w2"])
-
 
 if __name__ == "__main__":
     unittest.main()

tests/test_distributed_value.py

@@ -13,7 +13,7 @@ class DistributedValueTests(unittest.TestCase):
         from pathlib import Path
         from unittest.mock import MagicMock
 
-        module_path = Path(__file__).resolve().parents[3] / "nodes" / "utilities.py"
+        module_path = Path(__file__).resolve().parents[1] / "nodes" / "utilities.py"
         pkg_name = "dv_test_pkg"
 
         for mod_name in list(sys.modules):
@@ -28,18 +28,6 @@ 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
@@ -49,16 +37,6 @@ 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/test_job_timeout.py

@@ -9,7 +9,7 @@ from pathlib import Path
 
 
 def _load_job_timeout_module():
-    module_path = Path(__file__).resolve().parents[3] / "upscale" / "job_timeout.py"
+    module_path = Path(__file__).resolve().parents[1] / "upscale" / "job_timeout.py"
     package_name = "dist_job_timeout_testpkg"
 
     for mod_name in list(sys.modules):

tests/test_network_helpers.py

@@ -6,7 +6,7 @@ from pathlib import Path
 
 
 def _load_network_module():
-    module_path = Path(__file__).resolve().parents[3] / "utils" / "network.py"
+    module_path = Path(__file__).resolve().parents[1] / "utils" / "network.py"
 
     package_name = "dist_utils_testpkg"
     package_module = types.ModuleType(package_name)
@@ -17,10 +17,6 @@ def _load_network_module():
     logging_module.debug_log = lambda *_args, **_kwargs: None
     sys.modules[f"{package_name}.logging"] = logging_module
 
-    exceptions_module = types.ModuleType(f"{package_name}.exceptions")
-    exceptions_module.DistributedError = type("DistributedError", (Exception,), {})
-    sys.modules[f"{package_name}.exceptions"] = exceptions_module
-
     server_module = types.ModuleType("server")
     server_module.PromptServer = types.SimpleNamespace(
         instance=types.SimpleNamespace(address="127.0.0.1", port=8188, loop=None)
@@ -94,14 +90,48 @@ class NetworkHelpersTests(unittest.TestCase):
             "https://master.example.com",
         )
 
+    def test_build_master_url_ignores_stale_saved_port_and_uses_runtime_port(self):
+        cfg = {"master": {"host": "192.168.68.56", "port": 8001}}
+        prompt_server = types.SimpleNamespace(address="127.0.0.1", port=8188)
+        self.assertEqual(
+            network.build_master_url(config=cfg, prompt_server_instance=prompt_server),
+            "http://192.168.68.56:8188",
+        )
+
+    def test_build_master_url_keeps_explicit_port_in_host(self):
+        cfg = {"master": {"host": "192.168.68.56:8001"}}
+        prompt_server = types.SimpleNamespace(address="127.0.0.1", port=8188)
+        self.assertEqual(
+            network.build_master_url(config=cfg, prompt_server_instance=prompt_server),
+            "http://192.168.68.56:8001",
+        )
+
     def test_build_master_url_falls_back_to_server_address(self):
-        cfg = {"master": {"host": ""}}
-        prompt_server = types.SimpleNamespace(address="0.0.0.0", port=8190)  # nosec B104 - wildcard bind normalization test
+        cfg = {"master": {"host": "", "port": 8001}}
+        prompt_server = types.SimpleNamespace(address="0.0.0.0", port=8190)
         self.assertEqual(
             network.build_master_url(config=cfg, prompt_server_instance=prompt_server),
             "http://127.0.0.1:8190",
         )
 
+    def test_build_master_callback_url_uses_loopback_for_local_worker(self):
+        cfg = {"master": {"host": "192.168.68.56"}}
+        prompt_server = types.SimpleNamespace(address="127.0.0.1", port=8001)
+        worker = {"id": "w1", "type": "local", "host": "localhost", "port": 8189}
+        self.assertEqual(
+            network.build_master_callback_url(worker, config=cfg, prompt_server_instance=prompt_server),
+            "http://127.0.0.1:8001",
+        )
+
+    def test_build_master_callback_url_keeps_public_master_url_for_remote_worker(self):
+        cfg = {"master": {"host": "192.168.68.56"}}
+        prompt_server = types.SimpleNamespace(address="127.0.0.1", port=8001)
+        worker = {"id": "w2", "type": "remote", "host": "192.168.68.99", "port": 8189}
+        self.assertEqual(
+            network.build_master_callback_url(worker, config=cfg, prompt_server_instance=prompt_server),
+            "http://192.168.68.56:8001",
+        )
+
 
 if __name__ == "__main__":
     unittest.main()

tests/test_payload_parsers.py

@@ -15,7 +15,7 @@ except ImportError:
 
 def _load_payload_parsers_module():
     # payload_parsers.py has no relative imports; only stdlib + PIL
-    module_path = Path(__file__).resolve().parents[3] / "upscale" / "payload_parsers.py"
+    module_path = Path(__file__).resolve().parents[1] / "upscale" / "payload_parsers.py"
     spec = importlib.util.spec_from_file_location("upscale_payload_parsers", module_path)
     module = importlib.util.module_from_spec(spec)
     assert spec is not None and spec.loader is not None
@@ -86,27 +86,25 @@ class ParseTilesFromFormTests(unittest.TestCase):
 
     # --- happy paths ---
 
-    def test_single_tile_returns_one_entry(self):
+    def test_single_tile_returns_image_and_metadata(self):
         tiles = pp._parse_tiles_from_form(_make_form(1))
         self.assertEqual(len(tiles), 1)
-
-    def test_multiple_tiles_all_returned(self):
-        tiles = pp._parse_tiles_from_form(_make_form(3))
-        self.assertEqual(len(tiles), 3)
-
-    def test_tile_image_is_pil_image(self):
-        tiles = pp._parse_tiles_from_form(_make_form(1))
-        self.assertIsInstance(tiles[0]["image"], PILImage.Image)
-
-    def test_tile_metadata_fields_are_parsed(self):
-        tiles = pp._parse_tiles_from_form(_make_form(1))
         tile = tiles[0]
+        self.assertIsInstance(tile["image"], PILImage.Image)
         self.assertEqual(tile["tile_idx"], 0)
         self.assertEqual(tile["x"], 0)
         self.assertEqual(tile["y"], 0)
         self.assertEqual(tile["extracted_width"], 64)
         self.assertEqual(tile["extracted_height"], 64)
 
+    def test_multiple_tiles_preserve_count_order_and_coordinates(self):
+        tiles = pp._parse_tiles_from_form(_make_form(3))
+        self.assertEqual(len(tiles), 3)
+        for i, tile in enumerate(tiles):
+            self.assertEqual(tile["tile_idx"], i)
+        self.assertEqual(tiles[1]["x"], 64)
+        self.assertEqual(tiles[2]["x"], 128)
+
     def test_padding_is_parsed_from_form(self):
         tiles = pp._parse_tiles_from_form(_make_form(1, padding=16))
         self.assertEqual(tiles[0]["padding"], 16)
@@ -138,16 +136,6 @@ class ParseTilesFromFormTests(unittest.TestCase):
         self.assertNotIn("batch_idx", tiles[0])
         self.assertNotIn("global_idx", tiles[0])
 
-    def test_tile_indices_match_metadata_order(self):
-        tiles = pp._parse_tiles_from_form(_make_form(3))
-        for i, tile in enumerate(tiles):
-            self.assertEqual(tile["tile_idx"], i)
-
-    def test_x_coordinates_reflect_metadata(self):
-        tiles = pp._parse_tiles_from_form(_make_form(3))
-        self.assertEqual(tiles[1]["x"], 64)
-        self.assertEqual(tiles[2]["x"], 128)
-
     # --- error cases ---
 
     def test_missing_tiles_metadata_raises_value_error(self):

tests/test_queue_request.py

@@ -4,7 +4,7 @@ from pathlib import Path
 
 
 def _load_queue_request_module():
-    module_path = Path(__file__).resolve().parents[3] / "api" / "queue_request.py"
+    module_path = Path(__file__).resolve().parents[1] / "api" / "queue_request.py"
     spec = importlib.util.spec_from_file_location("queue_request", module_path)
     module = importlib.util.module_from_spec(spec)
     assert spec is not None and spec.loader is not None
@@ -26,26 +26,27 @@ class QueueRequestPayloadTests(unittest.TestCase):
 
     def test_normalizes_enabled_worker_ids(self):
         payload_data = self._base_payload()
-        payload_data["enabled_worker_ids"] = [" worker-a ", "worker-b", "worker-a"]
+        payload_data["enabled_worker_ids"] = ["a", 2, 3]
         payload_data["delegate_master"] = True
         payload = parse_queue_request_payload(
             payload_data
         )
-        self.assertEqual(payload.enabled_worker_ids, ["worker-a", "worker-b"])
+        self.assertEqual(payload.enabled_worker_ids, ["a", "2", "3"])
         self.assertTrue(payload.delegate_master)
 
     def test_supports_legacy_workers_field(self):
         payload_data = self._base_payload()
         payload_data.pop("enabled_worker_ids", None)
-        payload_data["workers"] = [{"id": "w1"}, "w2", {"id": "w3"}, {"name": "no-id"}]
+        payload_data["workers"] = [{"id": "w1"}, "w2", {"id": 3}, {"name": "no-id"}]
         payload = parse_queue_request_payload(
             payload_data
         )
-        self.assertEqual(payload.enabled_worker_ids, ["w1", "w2", "w3"])
+        self.assertEqual(payload.enabled_worker_ids, ["w1", "w2", "3"])
 
-    def test_supports_workflow_prompt_fallback(self):
+    def test_supports_auto_prepare_prompt_fallback(self):
         payload_data = self._base_payload()
         payload_data.pop("prompt", None)
+        payload_data["auto_prepare"] = True
         payload_data["workflow"] = {
             "prompt": {
                 "10": {"class_type": "DistributedCollector"},
@@ -55,6 +56,7 @@ 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()
@@ -72,6 +74,10 @@ 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)
@@ -85,7 +91,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_is_ignored_for_backward_compat(self):
+    def test_auto_prepare_false_still_falls_back_to_workflow_prompt(self):
         payload_data = self._base_payload()
         payload_data.pop("prompt", None)
         payload_data["auto_prepare"] = False
@@ -94,6 +100,13 @@ 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()
@@ -107,12 +120,6 @@ 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/test_runpod_client.py

@@ -0,0 +1,167 @@
+import importlib.util
+import json
+import sys
+import unittest
+from pathlib import Path
+
+ROOT = Path(__file__).resolve().parents[1]
+
+def load_module():
+    spec = importlib.util.spec_from_file_location("dist_runpod_client", ROOT / "api/orchestration/runpod_client.py")
+    module = importlib.util.module_from_spec(spec)
+    sys.modules[spec.name] = module
+    spec.loader.exec_module(module)
+    return module
+
+client = load_module()
+
+
+def workflow_for_request_size(target_size, timeout_seconds):
+    """Build an ASCII workflow whose compact UTF-8 request body has exact size."""
+    workflow = {"1": {"class_type": "Text", "inputs": {"value": ""}}}
+
+    def body_size(candidate):
+        payload = {
+            "input": {"workflow": candidate, "images": []},
+            "policy": {
+                "executionTimeout": timeout_seconds * 1000,
+                "ttl": timeout_seconds * 1000,
+            },
+        }
+        return len(json.dumps(
+            payload, allow_nan=False, ensure_ascii=False, separators=(",", ":")
+        ).encode("utf-8"))
+
+    padding = target_size - body_size(workflow)
+    if padding < 0:
+        raise AssertionError("target is smaller than the request fixture")
+    workflow["1"]["inputs"]["value"] = "x" * padding
+    if body_size(workflow) != target_size:
+        raise AssertionError("request fixture did not reach the exact target size")
+    return workflow
+
+class Response:
+    def __init__(self, status=200, payload=None):
+        self.status, self.payload = status, payload
+    async def __aenter__(self):
+        return self
+    async def __aexit__(self, *args):
+        return False
+    async def json(self):
+        return self.payload
+
+class Session:
+    def __init__(self, responses):
+        self.responses, self.calls = list(responses), []
+    def request(self, method, url, **kwargs):
+        self.calls.append((method, url, kwargs))
+        response = self.responses.pop(0)
+        if isinstance(response, Exception):
+            raise response
+        return response
+
+class ClientTests(unittest.IsolatedAsyncioTestCase):
+    async def test_submit_exact_generic_contract_and_once(self):
+        session = Session([Response(payload={"id": "job-1", "status": "IN_QUEUE"})])
+        api = client.RunpodClient(session, "endpoint123", "secret-value")
+        result = await api.submit({"1": {"class_type": "EmptyImage", "inputs": {}}}, [], 1800)
+        self.assertEqual(result, "job-1")
+        method, url, kwargs = session.calls[0]
+        self.assertEqual((method, url), ("POST", "https://api.runpod.ai/v2/endpoint123/run"))
+        self.assertIsInstance(kwargs["data"], bytes)
+        sent_payload = json.loads(kwargs["data"])
+        self.assertEqual(set(sent_payload["input"]), {"workflow", "images"})
+        self.assertEqual(sent_payload["policy"]["ttl"], 1800000)
+        self.assertFalse(kwargs["allow_redirects"])
+        self.assertEqual(kwargs["headers"]["Authorization"], "Bearer secret-value")
+
+    async def test_submit_timeout_is_ambiguous_and_never_retried(self):
+        session = Session([TimeoutError("secret-value in unsafe error")])
+        api = client.RunpodClient(session, "endpoint123", "secret-value")
+        with self.assertRaises(client.AmbiguousSubmission) as ctx:
+            await api.submit({}, [], 30)
+        self.assertNotIn("secret-value", str(ctx.exception))
+        self.assertEqual(len(session.calls), 1)
+
+    async def test_http_submit_failure_not_retried_or_leaked(self):
+        for status in [401, 429, 500, 302]:
+            session = Session([Response(status, {"error": "secret-value"})])
+            api = client.RunpodClient(session, "endpoint123", "secret-value")
+            with self.assertRaises(client.RunpodError) as ctx:
+                await api.submit({}, [], 30)
+            self.assertEqual(len(session.calls), 1)
+            self.assertNotIn("secret-value", str(ctx.exception))
+
+    async def test_health_zero_workers_is_valid(self):
+        session = Session([Response(payload={"workers": {"idle": 0}, "jobs": {"inQueue": 0}})])
+        api = client.RunpodClient(session, "endpoint123", "key")
+        self.assertEqual((await api.health())["workers"]["idle"], 0)
+
+    async def test_invalid_provider_id_no_network(self):
+        session = Session([])
+        with self.assertRaises(ValueError):
+            client.RunpodClient(session, "https://attacker.invalid", "key")
+        api = client.RunpodClient(session, "endpoint123", "key")
+        with self.assertRaises(ValueError):
+            await api.status("../other")
+        self.assertFalse(session.calls)
+
+    async def test_status_and_cancel_only_target_exact_job(self):
+        session = Session([Response(payload={"status": "IN_PROGRESS"}), Response(payload={"status": "CANCELLED"})])
+        api = client.RunpodClient(session, "endpoint123", "key")
+        await api.status("job-1")
+        await api.cancel("job-1")
+        self.assertEqual([c[:2] for c in session.calls], [("GET", "https://api.runpod.ai/v2/endpoint123/status/job-1"), ("POST", "https://api.runpod.ai/v2/endpoint123/cancel/job-1")])
+
+    async def test_missing_id_is_ambiguous(self):
+        api = client.RunpodClient(Session([Response(payload={"status": "IN_QUEUE"})]), "endpoint123", "key")
+        with self.assertRaises(client.AmbiguousSubmission):
+            await api.submit({}, [], 30)
+
+    async def test_large_payload_rejected_before_submit(self):
+        session = Session([])
+        api = client.RunpodClient(session, "endpoint123", "key")
+        with self.assertRaises(ValueError):
+            await api.submit({}, [{"image": "x" * client.MAX_REQUEST_BYTES}], 30)
+        self.assertFalse(session.calls)
+
+    async def test_exact_10mb_serialized_request_is_sent_as_validated_bytes(self):
+        session = Session([Response(payload={"id": "boundary-job"})])
+        api = client.RunpodClient(session, "endpoint123", "key")
+        timeout_seconds = 37
+        workflow = workflow_for_request_size(10_000_000, timeout_seconds)
+
+        result = await api.submit(workflow, [], timeout_seconds)
+
+        self.assertEqual(result, "boundary-job")
+        sent = session.calls[0][2]["data"]
+        self.assertIsInstance(sent, bytes)
+        self.assertEqual(len(sent), 10_000_000)
+        self.assertEqual(json.loads(sent)["policy"]["executionTimeout"], 37_000)
+
+    async def test_one_byte_over_10mb_is_rejected_without_request_or_payload_leak(self):
+        session = Session([Response(payload={"id": "must-not-be-used"})])
+        api = client.RunpodClient(session, "endpoint123", "key")
+        workflow = workflow_for_request_size(10_000_001, 37)
+        workflow["1"]["inputs"]["value"] = "private graph value" + workflow["1"]["inputs"]["value"][19:]
+
+        with self.assertRaisesRegex(ValueError, "10,000,001") as ctx:
+            await api.submit(workflow, [], 37)
+
+        self.assertNotIn("private graph value", str(ctx.exception))
+        self.assertFalse(session.calls)
+
+    async def test_wire_size_counts_utf8_graph_bytes_and_policy_fields(self):
+        session = Session([Response(payload={"id": "unicode-job"})])
+        api = client.RunpodClient(session, "endpoint123", "key")
+        workflow = {"1": {"class_type": "Text", "inputs": {"value": "café雪"}}}
+
+        await api.submit(workflow, [], 41)
+
+        sent = session.calls[0][2]["data"]
+        expected = json.dumps({
+            "input": {"workflow": workflow, "images": []},
+            "policy": {"executionTimeout": 41_000, "ttl": 41_000},
+        }, allow_nan=False, ensure_ascii=False, separators=(",", ":")).encode("utf-8")
+        self.assertEqual(sent, expected)
+        self.assertGreater(len(sent), len(sent.decode("utf-8")))

tests/test_serverless_collector.py

@@ -0,0 +1,202 @@
+import asyncio
+
+import pytest
+import torch
+
+from test_collector_list_inputs import _load_collector_module
+
+
+class _FakeManager:
+    def __init__(self, *, pending=False, owned=True):
+        self.pending = pending
+        self.owned = owned
+        self.pending_calls = []
+        self.owned_calls = []
+        self.cancel_calls = []
+
+    def is_pending(self, collector_job_id, participant_id):
+        self.pending_calls.append((collector_job_id, participant_id))
+        return self.pending
+
+    def owns(self, collector_job_id, participant_id):
+        self.owned_calls.append((collector_job_id, participant_id))
+        return self.owned
+
+    async def cancel(self, worker_id=None, execution_id=None, collector_job_id=None):
+        self.cancel_calls.append({
+            "worker_id": worker_id,
+            "execution_id": execution_id,
+            "collector_job_id": collector_job_id,
+        })
+        return 1
+
+
+def _set_up_collector(module, job_id, queue, manager=None):
+    module.prompt_server.distributed_jobs_lock = asyncio.Lock()
+    module.prompt_server.distributed_pending_jobs = {job_id: queue}
+    if manager is not None:
+        module.prompt_server.distributed_serverless = manager
+    return module.DistributedCollectorNode()
+
+
+def test_serverless_pending_waits_past_normal_worker_timeout():
+    module = _load_collector_module()
+    job_id = "collector-pending"
+    participant = "serverless:pool:1"
+    queue = asyncio.Queue()
+    manager = _FakeManager(pending=True)
+    collector = _set_up_collector(module, job_id, queue, manager)
+    module.get_worker_timeout_seconds = lambda: 0.01
+    module.HEARTBEAT_INTERVAL = 0.1
+    probe_calls = []
+    module.probe_worker = lambda *args, **kwargs: probe_calls.append(args)
+
+    async def deliver_after_provider_work():
+        await asyncio.sleep(0.14)
+        queue.put_nowait({
+            "worker_id": participant,
+            "image_index": 0,
+            "tensor": torch.ones(1, 2, 2, 2),
+            "is_last": True,
+        })
+
+    async def run_collection():
+        producer = asyncio.create_task(deliver_after_provider_work())
+        result = await collector.execute(
+            images=None,
+            audio=None,
+            multi_job_id=job_id,
+            enabled_worker_ids=f'["{participant}"]',
+            delegate_only=True,
+        )
+        await producer
+        return result
+
+    images, _audio = asyncio.run(run_collection())
+
+    assert tuple(images.shape) == (1, 2, 2, 2)
+    assert manager.pending_calls
+    assert probe_calls == []
+    assert job_id not in module.prompt_server.distributed_pending_jobs
+
+
+def test_serverless_queue_error_raises_and_cancels_only_collector_binding():
+    module = _load_collector_module()
+    job_id = "collector-error"
+    queue = asyncio.Queue()
+    queue.put_nowait({
+        "worker_id": "serverless:pool:1",
+        "serverless_error": "Provider execution failed safely",
+        "is_last": True,
+    })
+    manager = _FakeManager()
+    collector = _set_up_collector(module, job_id, queue, manager)
+
+    with pytest.raises(RuntimeError, match="serverless.*Provider execution failed safely"):
+        asyncio.run(collector.execute(
+            images=None,
+            audio=None,
+            multi_job_id=job_id,
+            enabled_worker_ids='["serverless:pool:1"]',
+            delegate_only=True,
+        ))
+
+    assert manager.cancel_calls == [{
+        "worker_id": None,
+        "execution_id": None,
+        "collector_job_id": job_id,
+    }]
+    assert job_id not in module.prompt_server.distributed_pending_jobs
+
+
+def test_owned_serverless_slot_without_pending_result_fails_instead_of_partial_success():
+    module = _load_collector_module()
+    job_id = "collector-stopped"
+    queue = asyncio.Queue()
+    manager = _FakeManager(pending=False, owned=True)
+    collector = _set_up_collector(module, job_id, queue, manager)
+    module.get_worker_timeout_seconds = lambda: 0.01
+    module.HEARTBEAT_INTERVAL = 0.1
+    module.probe_worker = lambda *_args, **_kwargs: pytest.fail("serverless slots are not worker URLs")
+
+    with pytest.raises(RuntimeError, match="serverless.*stopped.*result|serverless.*result.*stopped"):
+        asyncio.run(collector.execute(
+            images=None,
+            audio=None,
+            multi_job_id=job_id,
+            enabled_worker_ids='["serverless:pool:1"]',
+            delegate_only=True,
+        ))
+
+    assert manager.owned_calls
+    assert manager.cancel_calls[0]["collector_job_id"] == job_id
+    assert job_id not in module.prompt_server.distributed_pending_jobs
+
+
+def test_interruption_cancels_serverless_job_and_removes_collector_queue():
+    module = _load_collector_module()
+    job_id = "collector-interrupted"
+    queue = asyncio.Queue()
+    manager = _FakeManager(pending=True)
+    collector = _set_up_collector(module, job_id, queue, manager)
+    interrupted = module.comfy.model_management.InterruptProcessingException("stop")
+    module.comfy.model_management.throw_exception_if_processing_interrupted = lambda: (_ for _ in ()).throw(interrupted)
+
+    with pytest.raises(module.comfy.model_management.InterruptProcessingException, match="stop"):
+        asyncio.run(collector.execute(
+            images=None,
+            audio=None,
+            multi_job_id=job_id,
+            enabled_worker_ids='["serverless:pool:1"]',
+            delegate_only=True,
+        ))
+
+    assert manager.cancel_calls[0]["collector_job_id"] == job_id
+    assert job_id not in module.prompt_server.distributed_pending_jobs
+
+
+def test_serverless_tensor_shape_mismatch_raises_instead_of_silent_fallback():
+    module = _load_collector_module()
+    job_id = "collector-shape-mismatch"
+    queue = asyncio.Queue()
+    queue.put_nowait({
+        "worker_id": "serverless:pool:1",
+        "image_index": 0,
+        "tensor": torch.ones(1, 3, 2, 2),
+        "is_last": True,
+    })
+    manager = _FakeManager()
+    collector = _set_up_collector(module, job_id, queue, manager)
+    local_images = torch.zeros(1, 2, 2, 2)
+
+    with pytest.raises(RuntimeError, match="serverless.*shape|shape.*serverless"):
+        asyncio.run(collector.execute(
+            images=local_images,
+            audio=None,
+            multi_job_id=job_id,
+            enabled_worker_ids='["serverless:pool:1"]',
+        ))
+
+    assert manager.cancel_calls[0]["collector_job_id"] == job_id
+    assert job_id not in module.prompt_server.distributed_pending_jobs
+
+
+def test_regular_worker_timeout_keeps_partial_behavior_without_creating_manager():
+    module = _load_collector_module()
+    job_id = "collector-regular-timeout"
+    queue = asyncio.Queue()
+    collector = _set_up_collector(module, job_id, queue)
+    module.get_worker_timeout_seconds = lambda: 0.01
+    module.HEARTBEAT_INTERVAL = 0.1
+    local_images = torch.zeros(1, 2, 2, 2)
+    collector._get_serverless_manager = lambda *_args, **_kwargs: pytest.fail("regular collection must not create manager")
+
+    images, _audio = asyncio.run(collector.execute(
+        images=local_images,
+        audio=None,
+        multi_job_id=job_id,
+        enabled_worker_ids='["worker-a"]',
+    ))
+
+    assert torch.equal(images, local_images)
+    assert job_id not in module.prompt_server.distributed_pending_jobs

tests/test_serverless_config.py

@@ -0,0 +1,410 @@
+import importlib.util
+import json
+import os
+import stat
+import sys
+import tempfile
+import threading
+import types
+import unittest
+from pathlib import Path
+from unittest import mock
+
+
+def _load_serverless_config_module(*, without_fcntl=False):
+    module_path = Path(__file__).resolve().parents[1] / "utils" / "serverless_config.py"
+    module_name = (
+        "dist_serverless_config_no_fcntl_test_module"
+        if without_fcntl
+        else "dist_serverless_config_test_module"
+    )
+    package_name = module_name + "_pkg"
+    package = types.ModuleType(package_name)
+    package.__path__ = [str(module_path.parent)]
+    sys.modules[package_name] = package
+    module_name = package_name + ".serverless_config"
+    sys.modules.pop(module_name, None)
+    spec = importlib.util.spec_from_file_location(module_name, module_path)
+    assert spec is not None and spec.loader is not None
+    module = importlib.util.module_from_spec(spec)
+    sys.modules[module_name] = module
+    if without_fcntl:
+        with mock.patch.dict(sys.modules, {"fcntl": None}):
+            spec.loader.exec_module(module)
+    else:
+        spec.loader.exec_module(module)
+    return module
+
+
+serverless_config = _load_serverless_config_module()
+
+
+class ValidateServerlessWorkerTests(unittest.TestCase):
+    def test_returns_copy_with_runpod_defaults(self):
+        worker = {"id": "worker-1", "endpoint_id": "ep_123-abc", "name": "GPU"}
+
+        result = serverless_config.validate_serverless_worker(worker)
+
+        self.assertEqual(
+            result,
+            {
+                "id": "worker-1",
+                "endpoint_id": "ep_123-abc",
+                "name": "GPU",
+                "type": "serverless",
+                "provider": "runpod",
+                "jobs_per_run": 1,
+                "timeout_seconds": 1800,
+            },
+        )
+        self.assertIsNot(result, worker)
+        self.assertEqual(worker, {"id": "worker-1", "endpoint_id": "ep_123-abc", "name": "GPU"})
+
+    def test_accepts_inclusive_numeric_bounds(self):
+        for jobs, timeout in ((1, 30), (16, 86400)):
+            with self.subTest(jobs=jobs, timeout=timeout):
+                result = serverless_config.validate_serverless_worker(
+                    {
+                        "id": "worker-1",
+                        "endpoint_id": "ep-123",
+                        "jobs_per_run": jobs,
+                        "timeout_seconds": timeout,
+                    }
+                )
+                self.assertEqual(result["jobs_per_run"], jobs)
+                self.assertEqual(result["timeout_seconds"], timeout)
+
+    def test_legacy_concurrency_is_accepted_and_dropped(self):
+        # Configs saved before Images per run replaced the separate setting.
+        for legacy in (1, 2, 16, 17, True, "3"):
+            with self.subTest(legacy=legacy):
+                result = serverless_config.validate_serverless_worker(
+                    {"id": "worker-1", "endpoint_id": "ep-123", "jobs_per_run": 2, "concurrency": legacy}
+                )
+                self.assertNotIn("concurrency", result)
+                self.assertEqual(result["jobs_per_run"], 2)
+
+    def test_endpoint_url_is_rejected_with_a_plain_message(self):
+        with self.assertRaisesRegex(ValueError, "not a URL"):
+            serverless_config.validate_serverless_worker(
+                {"id": "worker-1", "endpoint_id": "https://api.runpod.ai/v2/abc123/run"}
+            )
+
+    def test_rejects_missing_or_empty_id_and_endpoint(self):
+        base: dict[str, object] = {"id": "worker-1", "endpoint_id": "ep-123"}
+        for field, value in (("id", None), ("id", "  "), ("endpoint_id", ""), ("endpoint_id", "https://api.runpod.io/ep-123")):
+            with self.subTest(field=field, value=value):
+                worker = dict(base)
+                worker[field] = value
+                with self.assertRaises(ValueError):
+                    serverless_config.validate_serverless_worker(worker)
+
+        worker = dict(base)
+        del worker["id"]
+        with self.assertRaises(ValueError):
+            serverless_config.validate_serverless_worker(worker)
+
+    def test_endpoint_id_must_be_ascii_opaque_identifier(self):
+        for endpoint_id in ("ep/123", "ep 123", "ép-123", "ep:123", "https://example.test"):
+            with self.subTest(endpoint_id=endpoint_id):
+                with self.assertRaises(ValueError):
+                    serverless_config.validate_serverless_worker(
+                        {"id": "worker-1", "endpoint_id": endpoint_id}
+                    )
+
+    def test_endpoint_id_matches_runpod_client_length_limit(self):
+        valid_id = "e" * 200
+        result = serverless_config.validate_serverless_worker(
+            {"id": "worker-1", "endpoint_id": valid_id}
+        )
+        self.assertEqual(result["endpoint_id"], valid_id)
+        with self.assertRaises(ValueError):
+            serverless_config.validate_serverless_worker(
+                {"id": "worker-1", "endpoint_id": "e" * 201}
+            )
+
+    def test_enabled_must_be_boolean_when_supplied(self):
+        base = {"id": "worker-1", "endpoint_id": "ep-123"}
+        for enabled in (True, False):
+            with self.subTest(enabled=enabled):
+                self.assertIs(
+                    serverless_config.validate_serverless_worker(
+                        {**base, "enabled": enabled}
+                    )["enabled"],
+                    enabled,
+                )
+        for enabled in ("false", 0, 1, None):
+            with self.subTest(enabled=enabled):
+                with self.assertRaises(ValueError):
+                    serverless_config.validate_serverless_worker(
+                        {**base, "enabled": enabled}
+                    )
+
+    def test_rejects_unknown_provider_and_non_serverless_type(self):
+        base = {"id": "worker-1", "endpoint_id": "ep-123"}
+        for extra in ({"provider": "other"}, {"type": "local"}):
+            with self.subTest(extra=extra):
+                with self.assertRaises(ValueError):
+                    serverless_config.validate_serverless_worker({**base, **extra})
+
+    def test_numeric_fields_reject_bool_float_and_out_of_range(self):
+        base = {"id": "worker-1", "endpoint_id": "ep-123"}
+        invalid_values = {
+            "jobs_per_run": (True, 1.0, 0, 17),
+            "timeout_seconds": (True, 30.0, 29, 86401),
+        }
+        for field, values in invalid_values.items():
+            for value in values:
+                with self.subTest(field=field, value=value):
+                    with self.assertRaises(ValueError):
+                        serverless_config.validate_serverless_worker({**base, field: value})
+
+    def test_requires_worker_object(self):
+        for worker in (None, [], "worker"):
+            with self.subTest(worker=worker):
+                with self.assertRaises(ValueError):
+                    serverless_config.validate_serverless_worker(worker)
+
+
+class SecretStoreTests(unittest.TestCase):
+    def setUp(self):
+        self.tempdir = tempfile.TemporaryDirectory()
+        self.root = Path(self.tempdir.name)
+        self.path = self.root / "private" / "runpod-secrets.json"
+
+    def tearDown(self):
+        self.tempdir.cleanup()
+
+    def test_set_get_has_delete_and_persists_as_json_mapping(self):
+        store = serverless_config.SecretStore(self.path)
+        store.set("worker-1", "secret-value")
+
+        self.assertEqual(store.get("worker-1"), "secret-value")
+        self.assertTrue(store.has("worker-1"))
+        self.assertIsNone(store.get("missing"))
+        self.assertFalse(store.has("missing"))
+        self.assertEqual(json.loads(self.path.read_text(encoding="utf-8")), {"worker-1": "secret-value"})
+        self.assertEqual(serverless_config.SecretStore(self.path).get("worker-1"), "secret-value")
+
+        store.delete("worker-1")
+        self.assertIsNone(store.get("worker-1"))
+        self.assertFalse(store.has("worker-1"))
+
+    def test_rejects_api_keys_with_newlines_or_over_4096_characters(self):
+        store = serverless_config.SecretStore(self.path)
+        for key in ("key\nwith-line-feed", "key\rwith-carriage-return", "k" * 4097):
+            with self.subTest(key_length=len(key)):
+                with self.assertRaises(serverless_config.SecretStoreError):
+                    store.set("worker-1", key)
+        store.set("worker-1", "k" * 4096)
+        self.assertEqual(store.get("worker-1"), "k" * 4096)
+
+    def test_new_secret_file_and_directory_have_restrictive_modes(self):
+        store = serverless_config.SecretStore(self.path)
+        store.set("worker-1", "secret-value")
+
+        self.assertEqual(stat.S_IMODE(self.path.stat().st_mode), 0o600)
+        self.assertEqual(stat.S_IMODE(self.path.parent.stat().st_mode), 0o700)
+
+    def test_missing_store_reads_as_empty_without_creating_files(self):
+        store = serverless_config.SecretStore(self.path)
+
+        self.assertIsNone(store.get("worker-1"))
+        self.assertFalse(store.has("worker-1"))
+        self.assertFalse(self.path.parent.exists())
+
+    def test_corrupt_store_fails_closed_without_overwriting_or_leaking_contents(self):
+        self.path.parent.mkdir()
+        corrupt = b'{"partial": "NEVER-LEAK-THIS-secret'
+        self.path.write_bytes(corrupt)
+        store = serverless_config.SecretStore(self.path)
+
+        for operation in (
+            lambda: store.get("worker-1"),
+            lambda: store.set("worker-1", "replacement-secret"),
+            lambda: store.delete("worker-1"),
+        ):
+            with self.subTest(operation=operation):
+                with self.assertRaises(Exception) as raised:
+                    operation()
+                self.assertNotIn("NEVER-LEAK-THIS", str(raised.exception))
+                self.assertNotIn("replacement-secret", str(raised.exception))
+                self.assertEqual(self.path.read_bytes(), corrupt)
+
+    def test_rejects_non_mapping_or_non_string_store_values_without_leaking_data(self):
+        for corrupt in (b'["opaque-stored-value"]', b'{"worker": 23}', b'{"": "opaque-stored-value"}'):
+            with self.subTest(corrupt=corrupt):
+                self.path.parent.mkdir(exist_ok=True)
+                self.path.write_bytes(corrupt)
+                with self.assertRaises(Exception) as raised:
+                    serverless_config.SecretStore(self.path).get("worker")
+                self.assertNotIn("opaque-stored-value", str(raised.exception))
+
+    def test_rejects_symlink_file_and_preserves_link_target(self):
+        target = self.root / "target.json"
+        target.write_text('{"worker-1": "target-secret"}', encoding="utf-8")
+        self.path.parent.mkdir()
+        self.path.symlink_to(target)
+        store = serverless_config.SecretStore(self.path)
+
+        with self.assertRaises(Exception):
+            store.get("worker-1")
+        with self.assertRaises(Exception):
+            store.set("worker-1", "replacement-secret")
+        self.assertEqual(target.read_text(encoding="utf-8"), '{"worker-1": "target-secret"}')
+
+    def test_rejects_symlinked_parent_directory(self):
+        real_dir = self.root / "real"
+        real_dir.mkdir()
+        link_dir = self.root / "linked"
+        link_dir.symlink_to(real_dir, target_is_directory=True)
+        store = serverless_config.SecretStore(link_dir / "secrets.json")
+
+        with self.assertRaises(Exception):
+            store.set("worker-1", "secret-value")
+        self.assertFalse((real_dir / "secrets.json").exists())
+
+    def test_rejects_parent_traversal_in_store_path(self):
+        store = serverless_config.SecretStore(self.root / "nested" / ".." / "secrets.json")
+
+        with self.assertRaises(Exception):
+            store.set("worker-1", "secret-value")
+        self.assertFalse((self.root / "secrets.json").exists())
+
+    def test_threaded_updates_from_independent_instances_are_not_lost(self):
+        workers = 12
+        barrier = threading.Barrier(workers)
+        errors = []
+
+        def write_one(index):
+            try:
+                barrier.wait(timeout=5)
+                serverless_config.SecretStore(self.path).set(f"worker-{index}", f"key-{index}")
+            except Exception as exc:  # collected to fail in the test thread
+                errors.append(exc)
+
+        threads = [threading.Thread(target=write_one, args=(i,)) for i in range(workers)]
+        for thread in threads:
+            thread.start()
+        for thread in threads:
+            thread.join(timeout=10)
+
+        self.assertFalse(any(thread.is_alive() for thread in threads), "writer thread did not finish")
+        self.assertEqual(errors, [])
+        saved = json.loads(self.path.read_text(encoding="utf-8"))
+        self.assertEqual(saved, {f"worker-{i}": f"key-{i}" for i in range(workers)})
+
+
+class PublicConfigTests(unittest.TestCase):
+    def test_redacts_credentials_preserves_other_fields_and_marks_only_serverless(self):
+        with tempfile.TemporaryDirectory() as tempdir:
+            store = serverless_config.SecretStore(Path(tempdir) / "private" / "keys.json")
+            store.set("serverless-1", "stored-api-key")
+            config = {
+                "workers": [
+                    {
+                        "id": "serverless-1",
+                        "type": "serverless",
+                        "provider": "runpod",
+                        "endpoint_id": "ep-123",
+                        "api_key": "inline-secret",
+                        "clear_api_key": "another-secret",
+                        "credential": "credential-secret",
+                        "nested": {"runpod_credential": "nested-secret", "safe": "kept"},
+                        "name": "remote GPU",
+                    },
+                    {"id": "serverless-2", "type": "serverless", "name": "not configured"},
+                    {"id": "local-1", "type": "local", "api_key": "local-secret"},
+                ],
+                "settings": {"mode": "preserved", "credentials_note": "remove-me"},
+            }
+            original = json.loads(json.dumps(config))
+
+            result = serverless_config.public_config(config, store)
+
+        first, second, local = result["workers"]
+        self.assertEqual(first["api_key_configured"], True)
+        self.assertEqual(second["api_key_configured"], False)
+        self.assertNotIn("api_key_configured", local)
+        self.assertNotIn("api_key", first)
+        self.assertNotIn("clear_api_key", first)
+        self.assertNotIn("credential", first)
+        self.assertNotIn("runpod_credential", first["nested"])
+        self.assertEqual(first["nested"]["safe"], "kept")
+        self.assertNotIn("api_key", local)
+        self.assertNotIn("credentials_note", result["settings"])
+        self.assertEqual(first["name"], "remote GPU")
+        self.assertEqual(config, original)
+        self.assertIsNot(result, config)
+        self.assertIsNot(result["workers"], config["workers"])
+
+    def test_redacts_normalized_secret_field_aliases(self):
+        config = {
+            "workers": [{
+                "id": "serverless-1",
+                "type": "serverless",
+                "apiKey": "camel-api-key",
+                "APIKEY": "upper-api-key",
+                "clearApiKey": "camel-clear-key",
+                "CLEAR_API_KEY": "upper-clear-key",
+                "runpod_CREDENTIAL": "provider-credential",
+                "nested": {"credential_note": "nested-credential"},
+                "safe": "preserved",
+            }]
+        }
+
+        result = serverless_config.public_config(
+            config, serverless_config.SecretStore(self.root_path())
+        )
+
+        worker = result["workers"][0]
+        self.assertEqual(worker, {
+            "id": "serverless-1",
+            "type": "serverless",
+            "nested": {},
+            "safe": "preserved",
+            "api_key_configured": False,
+        })
+
+    def test_handles_non_dict_config_without_mutating_it(self):
+        config = ["safe", {"api_key": "secret"}]
+        result = serverless_config.public_config(config, serverless_config.SecretStore(self.root_path()))
+        self.assertEqual(result, ["safe", {}])
+        self.assertEqual(config, ["safe", {"api_key": "secret"}])
+
+    def test_module_import_and_non_serverless_redaction_when_fcntl_unavailable(self):
+        module = _load_serverless_config_module(without_fcntl=True)
+        store = module.SecretStore(self.root_path())
+        config = {
+            "workers": [
+                {"id": "local-1", "type": "local", "api_key": "inline-secret"}
+            ]
+        }
+
+        result = module.public_config(config, store)
+
+        self.assertEqual(result, {"workers": [{"id": "local-1", "type": "local"}]})
+        self.assertEqual(config["workers"][0]["api_key"], "inline-secret")
+
+    def test_secret_operations_fail_clearly_without_fcntl(self):
+        module = _load_serverless_config_module(without_fcntl=True)
+        store = module.SecretStore(self.root_path())
+        unsupported_message = "Serverless credential storage is currently supported on POSIX masters only"
+
+        for operation in (
+            lambda: store.get("worker-1"),
+            lambda: store.set("worker-1", "api-key"),
+            lambda: store.delete("worker-1"),
+            lambda: store.has("worker-1"),
+        ):
+            with self.subTest(operation=operation):
+                with self.assertRaisesRegex(module.SecretStoreError, unsupported_message):
+                    operation()
+
+    @staticmethod
+    def root_path():
+        return Path(tempfile.gettempdir()) / "not-created-serverless-secret-store" / "secrets.json"
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/test_serverless_integration.py

@@ -0,0 +1,183 @@
+"""Offline integration: real graph preparation + orchestration, fake provider."""
+import asyncio
+import aiohttp as real_aiohttp
+import importlib.util
+import json
+import sys
+import types
+import unittest
+from pathlib import Path
+from unittest.mock import AsyncMock, patch
+
+ROOT = Path(__file__).resolve().parents[1]
+PKG = "dist_serverless_integration"
+
+
+def module(name, **attrs):
+    value = types.ModuleType(name)
+    value.__dict__.update(attrs)
+    return value
+
+
+def load_modules():
+    modules = {}
+    for part, directory in [("", ROOT), (".api", ROOT / "api"), (".api.orchestration", ROOT / "api/orchestration"), (".utils", ROOT / "utils")]:
+        modules[PKG + part] = module(PKG + part, __path__=[str(directory)])
+    ps = types.SimpleNamespace(routes=types.SimpleNamespace(get=lambda _: lambda fn: fn, post=lambda _: lambda fn: fn))
+    modules["server"] = module("server", PromptServer=types.SimpleNamespace(instance=ps))
+    modules[PKG + ".utils.logging"] = module("logging", debug_log=lambda *a: None, log=lambda *a: None)
+    modules[PKG + ".utils.trace_logger"] = module("trace_logger", trace_debug=lambda *a: None)
+    modules[PKG + ".utils.config"] = module("config", load_config=lambda: {})
+    modules[PKG + ".utils.async_helpers"] = module("async_helpers", queue_prompt_payload=AsyncMock())
+    modules[PKG + ".utils.network"] = module("network", build_master_url=lambda **k: "http://master", build_master_callback_url=lambda *a, **k: "http://master")
+    modules[PKG + ".api.orchestration.dispatch"] = module("dispatch", dispatch_worker_prompt=AsyncMock(), select_active_workers=AsyncMock(), select_least_busy_worker=AsyncMock())
+    modules[PKG + ".api.orchestration.media_sync"] = module("media_sync", convert_paths_for_platform=lambda p, _: p, fetch_worker_path_separator=AsyncMock(), fetch_worker_has_node=AsyncMock(return_value=False), sync_worker_media=AsyncMock())
+    with patch.dict(sys.modules, modules):
+        spec = importlib.util.spec_from_file_location(PKG + ".api.queue_orchestration", ROOT / "api/queue_orchestration.py")
+        q = importlib.util.module_from_spec(spec)
+        spec.loader.exec_module(q)
+        # Keep imported package namespace only; restore shared Comfy runtime stub.
+        modules.update({k: v for k, v in sys.modules.items() if k.startswith(PKG)})
+    return q, modules, ps
+
+
+class IntegrationTests(unittest.IsolatedAsyncioTestCase):
+    async def asyncSetUp(self):
+        self.q, modules, self.ps = load_modules()
+        self.patch = patch.dict(sys.modules, modules)
+        self.patch.start()
+        self.addCleanup(self.patch.stop)
+        self.pool = {"id": "pool", "name": "Runpod", "type": "serverless", "provider": "runpod", "endpoint_id": "endpoint", "enabled": True, "jobs_per_run": 2, "concurrency": 1, "timeout_seconds": 30}
+        self.q.load_config = lambda: {"workers": [self.pool], "settings": {"master_delegate_only": True}}
+        self.q.select_active_workers = AsyncMock(side_effect=lambda workers, _ws, delegate, **kw: (workers, delegate))
+        self.q.queue_prompt_payload = AsyncMock(return_value={"prompt_id": "master-prompt", "number": 1})
+        self.manager = types.SimpleNamespace(start=AsyncMock(), cancel=AsyncMock(), ensure_capacity=lambda _: None)
+        sys.modules[PKG + ".api.orchestration.serverless_runtime"] = module("runtime", get_manager=lambda: self.manager)
+        sys.modules[PKG + ".api.orchestration.serverless_media"] = module("media", prepare_images=AsyncMock(return_value=[]))
+        self.prompt = {
+            "1": {"class_type": "EmptyImage", "inputs": {"width": 32, "height": 32, "batch_size": 1, "color": 0}},
+            "2": {"class_type": "DistributedCollector", "inputs": {"images": ["1", 0], "load_balance": False}},
+            "3": {"class_type": "SaveImage", "inputs": {"images": ["2", 0], "filename_prefix": "local"}},
+        }
+
+    async def test_real_compiler_slots_to_master_collector_without_callbacks(self):
+        result = await self.q.orchestrate_distributed_execution(self.prompt, None, "client")
+        self.assertEqual(result[2], 2)
+        self.q.dispatch_worker_prompt.assert_not_awaited()
+        self.manager.start.assert_awaited_once()
+        jobs = self.manager.start.call_args.args[2]
+        self.assertEqual([job["worker"]["id"] for job in jobs], ["serverless:pool:1", "serverless:pool:2"])
+        self.assertEqual(len({job["manifest"][0]["prefix"] for job in jobs}), 2)
+        for job in jobs:
+            graph = job["workflow"]
+            self.assertTrue(any(n["class_type"] == "SaveImage" for n in graph.values()))
+            self.assertFalse(any("Distributed" in n["class_type"] for n in graph.values()))
+            self.assertNotIn("http://master", json.dumps(graph))
+        master_graph = self.q.queue_prompt_payload.call_args.args[0]
+        self.assertEqual(json.loads(master_graph["2"]["inputs"]["enabled_worker_ids"]), ["serverless:pool:1", "serverless:pool:2"])
+
+    async def test_compile_pack_provider_decode_and_collect_real_components(self):
+        """Exercise the whole image path, mocking only the provider and Comfy queue."""
+        import base64
+        import io
+        import tempfile
+        import torch
+        from PIL import Image
+        from test_collector_list_inputs import _load_collector_module
+
+        sys.modules["aiohttp"] = real_aiohttp
+        for leaf in ("serverless_media", "serverless_runtime"):
+            name = PKG + ".api.orchestration." + leaf
+            spec = importlib.util.spec_from_file_location(name, ROOT / "api/orchestration" / (leaf + ".py"))
+            real_module = importlib.util.module_from_spec(spec)
+            sys.modules[name] = real_module
+            spec.loader.exec_module(real_module)
+        runtime = sys.modules[PKG + ".api.orchestration.serverless_runtime"]
+        with tempfile.TemporaryDirectory() as directory:
+            root = Path(directory)
+            image_buffer = io.BytesIO()
+            Image.new("RGB", (2, 3), color=(255, 0, 128)).save(image_buffer, format="PNG")
+            image_bytes = image_buffer.getvalue()
+            (root / "source.png").write_bytes(image_bytes)
+            paths = module("folder_paths", get_input_directory=lambda: directory,
+                           get_output_directory=lambda: directory, get_temp_directory=lambda: directory,
+                           get_annotated_filepath=lambda name: str(root / name))
+            submissions = []
+
+            class FakeProvider:
+                async def submit(self, graph, images, timeout_seconds):
+                    self.prefix = next(n["inputs"]["filename_prefix"] for n in graph.values() if n["class_type"] == "SaveImage")
+                    submissions.append((graph, images))
+                    self_image = next(n["inputs"]["image"] for n in graph.values() if n["class_type"] == "LoadImage")
+                    assert self_image == images[0]["name"]
+                    assert base64.b64decode(images[0]["image"]) == image_bytes
+                    return "provider-job-" + str(len(submissions))
+
+                async def status(self, _job_id):
+                    filename = self.prefix + "_00001_.png"
+                    return {"status": "COMPLETED", "output": {"images": [{
+                        "filename": filename, "artifact_path": filename, "artifact_type": "output",
+                        "type": "base64", "data": base64.b64encode(image_bytes).decode(),
+                    }]}}
+
+            manager = runtime.ServerlessManager(self.ps, None, root / "state" / "jobs.json", None)
+            manager._client = AsyncMock(side_effect=lambda _: FakeProvider())
+            self.ps.distributed_serverless = manager
+            self.prompt["1"] = {"class_type": "LoadImage", "inputs": {"image": "source.png"}}
+            collector_module = _load_collector_module()
+            collector_module.prompt_server = self.ps
+            sys.modules["aiohttp"] = real_aiohttp
+            with patch.object(runtime, "get_manager", return_value=manager), patch.dict(sys.modules, {"folder_paths": paths}):
+                result = await self.q.orchestrate_distributed_execution(self.prompt, None, "client")
+                master_graph = self.q.queue_prompt_payload.call_args.args[0]
+                inputs = master_graph["2"]["inputs"]
+                images, _audio = await asyncio.wait_for(collector_module.DistributedCollectorNode().execute(
+                    images=None, audio=None, multi_job_id=inputs["multi_job_id"],
+                    enabled_worker_ids=inputs["enabled_worker_ids"], delegate_only=True,
+                ), timeout=3)
+                await asyncio.gather(*manager.tasks.values())
+            self.assertEqual(result[2], 2)
+            self.assertEqual(len(submissions), 2)
+            self.assertEqual(tuple(images.shape), (2, 3, 2, 3))
+            self.assertEqual(images.dtype, torch.float32)
+            self.assertTrue(torch.equal(images[0], images[1]))
+            self.assertEqual(float(images[0, 0, 0, 0]), 1.0)
+            self.assertEqual([record["status"] for record in manager.ledger.all()], ["COMPLETED", "COMPLETED"])
+            self.assertEqual(self.ps.distributed_pending_jobs, {})
+            self.q.dispatch_worker_prompt.assert_not_awaited()
+
+    async def test_audio_rejected_before_any_dispatch_or_master_queue(self):
+        self.prompt["2"]["inputs"]["audio"] = ["9", 0]
+        self.prompt["9"] = {"class_type": "LoadAudio", "inputs": {"audio": "test.wav"}}
+        with self.assertRaises(ValueError):
+            await self.q.orchestrate_distributed_execution(self.prompt, None, "client")
+        self.manager.start.assert_not_awaited()
+        self.q.dispatch_worker_prompt.assert_not_awaited()
+        self.q.queue_prompt_payload.assert_not_awaited()
+
+    async def test_disabled_pool_excluded(self):
+        self.pool["enabled"] = False
+        self.assertEqual(self.q._resolve_enabled_workers({"workers": [self.pool]}), [])
+
+    async def test_master_validation_failure_submits_no_paid_job(self):
+        self.q.queue_prompt_payload.side_effect = ValueError("invalid local graph")
+        with self.assertRaisesRegex(ValueError, "invalid local graph"):
+            await self.q.orchestrate_distributed_execution(self.prompt, None, "client")
+        self.manager.start.assert_not_awaited()
+
+    async def test_oversized_serverless_request_fails_before_dispatch_or_master_queue(self):
+        media_module = sys.modules[PKG + ".api.orchestration.serverless_media"]
+        media_module.prepare_images.return_value = [{
+            "name": "private-image-name.png",
+            "image": "x" * 10_000_001,
+        }]
+
+        with self.assertRaisesRegex(ValueError, "10,000,") as ctx:
+            await self.q.orchestrate_distributed_execution(self.prompt, None, "client")
+
+        self.assertNotIn("private-image-name.png", str(ctx.exception))
+        self.assertNotIn("x" * 100, str(ctx.exception))
+        self.q.dispatch_worker_prompt.assert_not_awaited()
+        self.q.queue_prompt_payload.assert_not_awaited()
+        self.manager.start.assert_not_awaited()
+        self.assertEqual(self.ps.distributed_pending_jobs, {})

tests/test_serverless_lifecycle.py

@@ -0,0 +1,198 @@
+"""Billing-safe lifecycle regressions: local completion is not remote completion."""
+import asyncio
+import tempfile
+import types
+import unittest
+from pathlib import Path
+from unittest.mock import AsyncMock, patch
+
+from test_serverless_runtime import runtime, Secrets
+
+
+class LifecycleTests(unittest.IsolatedAsyncioTestCase):
+    async def asyncSetUp(self):
+        self.temp = tempfile.TemporaryDirectory()
+        self.addCleanup(self.temp.cleanup)
+        self.server = types.SimpleNamespace(distributed_pending_jobs={"collector": asyncio.Queue()}, distributed_jobs_lock=asyncio.Lock())
+        self.path = Path(self.temp.name) / "jobs.json"
+        self.manager = runtime.ServerlessManager(self.server, Secrets(), self.path, AsyncMock(return_value=object()))
+        self.worker = {"id": "serverless:pool:1", "pool_id": "pool", "endpoint_id": "ep", "timeout_seconds": 30, "type": "serverless"}
+        self.prepared = {"worker": self.worker, "workflow": {}, "images": [], "manifest": [{"job_id": "collector", "prefix": "x", "kind": "image"}]}
+        self.api = types.SimpleNamespace(submit=AsyncMock(return_value="job"), status=AsyncMock(), cancel=AsyncMock(return_value={"status": "CANCELLED"}))
+        self.client_patch = patch.object(runtime, "RunpodClient", return_value=self.api)
+        self.client_patch.start()
+        self.addCleanup(self.client_patch.stop)
+        self.decode_patch = patch.object(runtime, "decode_results", AsyncMock(return_value={"collector": ["tensor"]}))
+        self.decode_patch.start()
+        self.addCleanup(self.decode_patch.stop)
+        # These regressions exercise capacity holding, so use a limit of one.
+        self.limit_patch = patch.object(runtime, "PER_WORKER_JOB_LIMIT", 1)
+        self.limit_patch.start()
+        self.addCleanup(self.limit_patch.stop)
+
+    async def asyncTearDown(self):
+        await self.manager.cancel()
+
+    def second(self, **changes):
+        return dict(self.prepared, worker=dict(self.worker, id="serverless:pool:2", **changes))
+
+    def retained(self, status="CANCEL_REQUESTED", job_id: str | None = "old-job"):
+        self.manager.ledger.update("old", worker_id="pool", endpoint_id="ep", job_id=job_id,
+                                   status=status, created_at=1, updated_at=1)
+
+    async def test_capacity_is_held_until_slow_cancellation_is_confirmed(self):
+        cancel_started, release_cancel = asyncio.Event(), asyncio.Event()
+        self.api.submit.side_effect = ["job-1", "job-2"]
+        self.api.status.side_effect = [runtime.RunpodError("poll failed"), {"status": "COMPLETED", "output": {}}]
+        async def cancel(_job_id):
+            cancel_started.set()
+            await release_cancel.wait()
+            return {"status": "CANCELLED"}
+        self.api.cancel.side_effect = cancel
+        await self.manager.start("exec", "prompt", [self.prepared, self.second()])
+        try:
+            await asyncio.wait_for(cancel_started.wait(), 1)
+            await asyncio.sleep(0.01)
+            self.assertEqual(self.api.submit.await_count, 1)
+        finally:
+            release_cancel.set()
+            await self.manager.wait_for_execution("exec")
+        self.assertEqual(self.api.submit.await_count, 2)
+
+    async def test_unconfirmed_cancel_holds_capacity_until_reconciled(self):
+        self.api.status.side_effect = runtime.RunpodError("poll failed")
+        self.api.cancel.return_value = {"status": "IN_PROGRESS"}
+        await self.manager.start("first", "prompt", [self.prepared])
+        await self.manager.wait_for_execution("first")
+        self.assertEqual(self.manager.ledger.all()[0]["status"], "CANCEL_REQUESTED")
+        self.assertTrue(self.manager.pool_is_active("pool"))
+        await self.manager.start("second", "prompt", [self.second()])
+        try:
+            await asyncio.sleep(0.01)
+            self.assertEqual(self.api.submit.await_count, 1)
+            self.api.status.side_effect = None
+            self.api.status.return_value = {"status": "COMPLETED", "output": {}}
+            await self.manager.reconcile(worker_id="pool")
+            await asyncio.wait_for(self.manager.wait_for_execution("second"), 1)
+            self.assertEqual(self.api.submit.await_count, 2)
+        finally:
+            self.api.cancel.return_value = {"status": "CANCELLED"}
+            await self.manager.cancel()
+
+    async def test_restarted_manager_reserves_known_unresolved_jobs(self):
+        self.retained("IN_PROGRESS")
+        self.manager = runtime.ServerlessManager(self.server, Secrets(), self.path, AsyncMock(return_value=object()))
+        self.assertTrue(self.manager.pool_is_active("pool"))
+        self.assertFalse(self.manager.pool_is_active("other"))
+        await self.manager.start("exec", "prompt", [self.second(timeout_seconds=0.02)])
+        await self.manager.wait_for_execution("exec")
+        self.api.submit.assert_not_awaited()
+        self.assertEqual(self.manager.ledger.all()[-1]["status"], "TIMED_OUT")
+
+    async def test_recovery_closes_never_submitted_waiting_records(self):
+        self.retained("WAITING", job_id=None)
+        self.manager = runtime.ServerlessManager(self.server, Secrets(), self.path, AsyncMock(return_value=object()))
+        self.assertFalse(self.manager.pool_is_active("pool"))
+        self.assertEqual(self.manager.ledger.all()[0]["status"], "CANCELLED")
+        self.api.submit.assert_not_awaited()
+
+    async def test_recovery_keeps_interrupted_submission_unresolved(self):
+        self.retained("SUBMITTING", job_id=None)
+        self.manager = runtime.ServerlessManager(self.server, Secrets(), self.path, AsyncMock(return_value=object()))
+        self.assertTrue(self.manager.pool_is_active("pool"))
+        self.assertEqual(self.manager.ledger.all()[0]["status"], "UNKNOWN")
+        await self.manager.reconcile(worker_id="pool")
+        self.api.status.assert_not_awaited()
+        self.api.submit.assert_not_awaited()
+
+    async def test_ambiguous_submission_requires_explicit_recovery_acknowledgment(self):
+        self.api.submit.side_effect = runtime.AmbiguousSubmission("acceptance unknown")
+        await self.manager.start("exec", "prompt", [self.prepared])
+        await self.manager.wait_for_execution("exec")
+        record = self.manager.ledger.all()[0]
+        self.assertTrue(self.manager.pool_is_active("pool"))
+        with self.assertRaises(ValueError):
+            await self.manager.acknowledge_recovery("pool", record["record_id"], provider_checked=False)
+        with self.assertRaises(ValueError):
+            await self.manager.acknowledge_recovery("other", record["record_id"], provider_checked=True)
+        await self.manager.acknowledge_recovery("pool", record["record_id"], provider_checked=True)
+        self.assertFalse(self.manager.pool_is_active("pool"))
+        self.assertEqual(self.manager.ledger.all()[0]["status"], "RESOLVED")
+        self.api.submit.assert_awaited_once()
+
+    async def test_cancel_already_completed_provider_job_is_terminal(self):
+        self.retained()
+        self.api.cancel.return_value = {"status": "COMPLETED"}
+        await self.manager.cancel(worker_id="pool")
+        self.assertEqual(self.manager.ledger.all()[0]["status"], "COMPLETED")
+        self.assertFalse(self.manager.pool_is_active("pool"))
+
+    async def test_reconciliation_failure_never_releases_capacity(self):
+        self.retained()
+        self.api.status.side_effect = runtime.RunpodError("not reachable")
+        await self.manager.reconcile(worker_id="pool")
+        self.assertTrue(self.manager.pool_is_active("pool"))
+        self.assertEqual(self.manager.ledger.all()[0]["status"], "CANCEL_REQUESTED")
+
+    async def test_cancel_before_first_task_step_finalizes_local_record(self):
+        await self.manager.start("exec", "prompt", [self.prepared])
+        await self.manager.cancel(worker_id="pool")
+        self.assertEqual(self.manager.ledger.all()[0]["status"], "CANCELLED")
+        self.assertFalse(self.server.distributed_pending_jobs["collector"].empty())
+        self.assertFalse(self.manager.pool_is_active("pool"))
+        self.api.submit.assert_not_awaited()
+
+    async def test_live_record_cannot_be_acknowledged_or_reconciled(self):
+        polling = asyncio.Event()
+        async def status(_job_id):
+            polling.set()
+            await asyncio.Event().wait()
+        self.api.status.side_effect = status
+        await self.manager.start("exec", "prompt", [self.prepared])
+        await asyncio.wait_for(polling.wait(), 1)
+        record = self.manager.ledger.all()[0]
+        with self.assertRaises(ValueError):
+            await self.manager.acknowledge_recovery("pool", record["record_id"], provider_checked=True)
+        await self.manager.reconcile(worker_id="pool")
+        self.assertEqual(self.api.status.await_count, 1)
+        self.assertTrue(self.manager.pool_is_active("pool"))
+
+    async def test_repeated_cancel_does_not_interrupt_provider_cleanup(self):
+        polling, cancelling, release = asyncio.Event(), asyncio.Event(), asyncio.Event()
+        async def status(_job_id):
+            polling.set()
+            await asyncio.Event().wait()
+        async def cancel(_job_id):
+            cancelling.set()
+            await release.wait()
+            return {"status": "CANCELLED"}
+        self.api.status.side_effect = status
+        self.api.cancel.side_effect = cancel
+        await self.manager.start("exec", "prompt", [self.prepared])
+        await polling.wait()
+        first = asyncio.create_task(self.manager.cancel(worker_id="pool"))
+        await cancelling.wait()
+        second = asyncio.create_task(self.manager.cancel(worker_id="pool"))
+        await asyncio.sleep(0)
+        release.set()
+        await asyncio.gather(first, second)
+        self.api.cancel.assert_awaited_once_with("job")
+        self.assertEqual(self.manager.ledger.all()[0]["status"], "CANCELLED")
+
+    async def test_partial_start_cleanup_does_not_leave_waiting_protection(self):
+        broken = self.second()
+        del broken["worker"]["endpoint_id"]
+        with self.assertRaises(KeyError):
+            await self.manager.start("exec", "prompt", [self.prepared, broken])
+        await self.manager.cancel(execution_id="exec")
+        self.assertFalse(self.manager.pool_is_active("pool"))
+        self.api.submit.assert_not_awaited()
+
+    async def test_public_jobs_never_hides_unresolved_recovery_ids(self):
+        for index in range(22):
+            self.manager.ledger.update(str(index), worker_id="pool", endpoint_id="ep", status="UNKNOWN", updated_at=index)
+        for index in range(5):
+            self.manager.ledger.update("done-" + str(index), worker_id="pool", status="COMPLETED", updated_at=30 + index)
+        records = self.manager.public_jobs("pool")
+        self.assertEqual({record["record_id"] for record in records}, {str(index) for index in range(22)})
+        self.assertTrue(all(record["status"] == "UNKNOWN" for record in records))

tests/test_serverless_media.py

@@ -0,0 +1,447 @@
+import base64
+import importlib.util
+import io
+import sys
+import tempfile
+import types
+import unittest
+from pathlib import Path
+from unittest.mock import AsyncMock, patch
+
+from PIL import Image
+
+ROOT = Path(__file__).resolve().parents[1]
+MODULE_PATH = ROOT / "api" / "orchestration" / "serverless_media.py"
+
+
+def load_module():
+    spec = importlib.util.spec_from_file_location("dist_serverless_media_test", MODULE_PATH)
+    assert spec is not None and spec.loader is not None
+    module = importlib.util.module_from_spec(spec)
+    sys.modules[spec.name] = module
+    spec.loader.exec_module(module)
+    return module
+
+
+media = load_module()
+
+
+def load_runpod_module():
+    spec = importlib.util.spec_from_file_location(
+        "dist_serverless_media_runpod_test", ROOT / "api" / "orchestration" / "runpod_client.py"
+    )
+    assert spec is not None and spec.loader is not None
+    module = importlib.util.module_from_spec(spec)
+    sys.modules[spec.name] = module
+    spec.loader.exec_module(module)
+    return module
+
+
+runpod = load_runpod_module()
+
+
+def png_bytes(mode="RGB", size=(2, 1), color=None):
+    if color is None:
+        color = (10, 20, 30, 40) if mode == "RGBA" else (10, 20, 30)
+    buffer = io.BytesIO()
+    Image.new(mode, size, color).save(buffer, format="PNG")
+    return buffer.getvalue()
+
+
+class FolderPathsFixture:
+    def __init__(self, root):
+        self.input = root / "input"
+        self.output = root / "output"
+        self.temp = root / "temp"
+        for directory in (self.input, self.output, self.temp):
+            directory.mkdir(parents=True, exist_ok=True)
+
+    def get_input_directory(self):
+        return str(self.input)
+
+    def get_output_directory(self):
+        return str(self.output)
+
+    def get_temp_directory(self):
+        return str(self.temp)
+
+    def get_annotated_filepath(self, filename):
+        value = filename.replace("\\", "/")
+        # Match real ComfyUI folder_paths: root annotations are suffixes,
+        # not directory prefixes (output/foo.png is an input subdirectory).
+        for marker, root in (("input", self.input), ("output", self.output), ("temp", self.temp)):
+            suffix = f" [{marker}]"
+            if value.endswith(suffix):
+                return str(root / value[:-len(suffix)])
+        return str(self.input / value)
+
+
+class ServerlessMediaInputTests(unittest.IsolatedAsyncioTestCase):
+    async def asyncSetUp(self):
+        self.tempdir = tempfile.TemporaryDirectory()
+        self.paths = FolderPathsFixture(Path(self.tempdir.name))
+        self.folder_paths = types.ModuleType("folder_paths")
+        self.folder_paths.get_input_directory = self.paths.get_input_directory
+        self.folder_paths.get_output_directory = self.paths.get_output_directory
+        self.folder_paths.get_temp_directory = self.paths.get_temp_directory
+        self.folder_paths.get_annotated_filepath = self.paths.get_annotated_filepath
+        self.folder_patch = patch.dict(sys.modules, {"folder_paths": self.folder_paths})
+        self.folder_patch.start()
+        self.addCleanup(self.folder_patch.stop)
+        self.addCleanup(self.tempdir.cleanup)
+
+    async def test_packs_image_and_rewrites_duplicate_loadimage_inputs(self):
+        payload = png_bytes()
+        (self.paths.input / "cat.png").write_bytes(payload)
+        workflow = {
+            "1": {"class_type": "LoadImage", "inputs": {"image": "cat.png [input]", "other": "unchanged"}},
+            "2": {"class_type": "LoadImage", "inputs": {"image": "cat.png [input]"}},
+            "3": {"class_type": "SaveImage", "inputs": {"filename_prefix": "stable"}},
+        }
+
+        packed = await media.prepare_images(workflow)
+
+        self.assertEqual(len(packed), 1)
+        self.assertEqual(set(packed[0]), {"name", "image"})
+        self.assertRegex(packed[0]["name"], r"^[A-Za-z0-9_.-]+\.png$")
+        self.assertEqual(base64.b64decode(packed[0]["image"], validate=True), payload)
+        self.assertEqual(workflow["1"]["inputs"]["image"], packed[0]["name"])
+        self.assertEqual(workflow["2"]["inputs"]["image"], packed[0]["name"])
+        self.assertEqual(workflow["1"]["inputs"]["other"], "unchanged")
+        self.assertEqual(workflow["3"]["inputs"]["filename_prefix"], "stable")
+
+    async def test_resolves_annotated_input_output_and_temp_sources(self):
+        sources = ((self.paths.input, "one.png"), (self.paths.output, "two.jpg"), (self.paths.temp, "three.webp"))
+        workflow = {}
+        for index, (directory, filename) in enumerate(sources, 1):
+            (directory / filename).write_bytes(png_bytes())
+            root = directory.name
+            workflow[str(index)] = {"class_type": "LoadImage", "inputs": {"image": f"{filename} [{root}]"}}
+
+        packed = await media.prepare_images(workflow)
+
+        self.assertEqual(len(packed), 3)
+        self.assertEqual({base64.b64decode(item["image"]) for item in packed}, {png_bytes()})
+        self.assertEqual({node["inputs"]["image"] for node in workflow.values()}, {item["name"] for item in packed})
+
+    async def test_directory_prefix_is_not_a_root_annotation(self):
+        (self.paths.input / "output").mkdir()
+        payload = png_bytes(color=(100, 0, 0))
+        (self.paths.input / "output" / "same.png").write_bytes(payload)
+        (self.paths.output / "same.png").write_bytes(png_bytes())
+        workflow = {"1": {"class_type": "LoadImage", "inputs": {"image": "output/same.png"}}}
+        packed = await media.prepare_images(workflow)
+        self.assertEqual(base64.b64decode(packed[0]["image"]), payload)
+
+    async def test_rejects_unknown_root_annotation(self):
+        (self.paths.input / "cat.png").write_bytes(png_bytes())
+        workflow = {"1": {"class_type": "LoadImage", "inputs": {"image": "cat.png [abc123]"}}}
+        with self.assertRaisesRegex(media.ServerlessMediaError, "annotation"):
+            await media.prepare_images(workflow)
+
+    async def test_recognizes_custom_file_image_reference(self):
+        (self.paths.input / "source.bmp").write_bytes(png_bytes())
+        workflow = {"1": {"class_type": "CustomImageLoader", "inputs": {"file": "source.bmp"}}}
+
+        packed = await media.prepare_images(workflow)
+
+        self.assertEqual(len(packed), 1)
+        self.assertEqual(workflow["1"]["inputs"]["file"], packed[0]["name"])
+
+    async def test_rejects_missing_and_dynamic_loadimage_values(self):
+        cases = (
+            {"1": {"class_type": "LoadImage", "inputs": {}}},
+            {"1": {"class_type": "LoadImage"}},
+            {"1": {"class_type": "LoadImage", "inputs": None}},
+            {"1": {"class_type": "LoadImage", "inputs": {"image": ["2", 0]}}},
+        )
+        for workflow in cases:
+            with self.subTest(workflow=workflow):
+                with self.assertRaisesRegex(media.ServerlessMediaError, "LoadImage"):
+                    await media.prepare_images(workflow)
+
+    async def test_rejects_missing_source_and_unapproved_extension(self):
+        for value in ("absent.png", "animation.gif", "clip.mp4"):
+            with self.subTest(value=value):
+                workflow = {"1": {"class_type": "LoadImage", "inputs": {"image": value}}}
+                with self.assertRaises(media.ServerlessMediaError):
+                    await media.prepare_images(workflow)
+
+    async def test_rejects_unsupported_audio_or_video_references_without_reading(self):
+        for key, value in (("audio", "private.wav"), ("video", "private.mp4"), ("file", "private.wav")):
+            with self.subTest(key=key):
+                workflow = {"1": {"class_type": "CustomLoader", "inputs": {key: value}}}
+                with self.assertRaisesRegex(media.ServerlessMediaError, "audio|video|media"):
+                    await media.prepare_images(workflow)
+
+    async def test_rejects_parent_traversal_and_absolute_outside_roots(self):
+        outside = Path(self.tempdir.name).parent / "outside.png"
+        outside.write_bytes(png_bytes())
+        self.addCleanup(lambda: outside.unlink(missing_ok=True))
+        for value in ("../outside.png", str(outside)):
+            with self.subTest(value=value):
+                workflow = {"1": {"class_type": "LoadImage", "inputs": {"image": value}}}
+                with self.assertRaises(media.ServerlessMediaError):
+                    await media.prepare_images(workflow)
+
+    async def test_rejects_symlink_escape_from_allowed_root(self):
+        outside = Path(self.tempdir.name).parent / "linked-outside.png"
+        outside.write_bytes(png_bytes())
+        self.addCleanup(lambda: outside.unlink(missing_ok=True))
+        (self.paths.input / "link.png").symlink_to(outside)
+        workflow = {"1": {"class_type": "LoadImage", "inputs": {"image": "link.png"}}}
+
+        with self.assertRaises(media.ServerlessMediaError):
+            await media.prepare_images(workflow)
+
+    async def test_rejects_individual_and_aggregate_input_size_overages(self):
+        huge = self.paths.input / "large.jpg"
+        with huge.open("wb") as handle:
+            handle.truncate(media.MAX_INPUT_IMAGE_BYTES + 1)
+        workflow = {"1": {"class_type": "LoadImage", "inputs": {"image": "large.jpg"}}}
+        with self.assertRaisesRegex(media.ServerlessMediaError, "size"):
+            await media.prepare_images(workflow)
+
+        huge.unlink()
+        first = self.paths.input / "first.png"
+        second = self.paths.input / "second.png"
+        first.write_bytes(b"a" * (media.MAX_INPUT_TOTAL_BYTES // 2 + 1))
+        second.write_bytes(b"b" * (media.MAX_INPUT_TOTAL_BYTES // 2 + 1))
+        workflow = {
+            "1": {"class_type": "LoadImage", "inputs": {"image": "first.png"}},
+            "2": {"class_type": "LoadImage", "inputs": {"image": "second.png"}},
+        }
+        with self.assertRaisesRegex(media.ServerlessMediaError, "size"):
+            await media.prepare_images(workflow)
+
+
+    async def test_accepts_raw_image_over_old_7mib_when_complete_request_fits_10mb(self):
+        raw_size = 7 * 1024 * 1024 + 1
+        source = self.paths.input / "large-input.png"
+        source.write_bytes(b"x" * raw_size)
+        workflow = {"1": {"class_type": "LoadImage", "inputs": {"image": "large-input.png"}}}
+
+        packed = await media.prepare_images(workflow)
+        request_body = runpod.serialize_request(workflow, packed, 30)
+
+        self.assertGreater(raw_size, 7 * 1024 * 1024)
+        self.assertEqual(len(base64.b64decode(packed[0]["image"])), raw_size)
+        self.assertLessEqual(len(request_body), 10_000_000)
+
+
+class ServerlessMediaOutputTests(unittest.IsolatedAsyncioTestCase):
+    def manifest(self, prefix="render"):
+        return [{"collector_id": "collector-1", "job_id": "job-1", "node_id": "12", "kind": "image", "prefix": prefix}]
+
+    def artifact(self, filename, payload=None, **extra):
+        entry = {
+            "filename": filename,
+            "artifact_type": "output",
+            "artifact_path": f"/safe/output/{filename}",
+            "subfolder": "",
+            "type": "base64",
+            "data": base64.b64encode(payload if payload is not None else png_bytes()).decode("ascii"),
+        }
+        entry.update(extra)
+        return entry
+
+    async def test_decodes_rgb_to_normalized_float32_and_orders_by_counter(self):
+        output = {"images": [self.artifact("render_00009_.png"), self.artifact("render_00002_.png")]}
+
+        result = await media.collect_images(output, self.manifest())
+
+        self.assertEqual(list(result), ["job-1"])
+        self.assertEqual(len(result["job-1"]), 2)
+        self.assertTrue(all(image.dtype.name == "float32" for image in result["job-1"]))
+        self.assertEqual(result["job-1"][0].shape, (1, 2, 3))
+        self.assertAlmostEqual(float(result["job-1"][0][0, 0, 0]), 10 / 255)
+
+    async def test_groups_images_under_manifest_job_ids(self):
+        manifest = [
+            self.manifest("first")[0],
+            {"collector_id": "collector-2", "job_id": "job-2", "node_id": "14", "kind": "image", "prefix": "second"},
+        ]
+        output = {"images": [self.artifact("first_00001_.png"), self.artifact("second_00001_.png")]}
+
+        result = await media.collect_images(output, manifest)
+
+        self.assertEqual(set(result), {"job-1", "job-2"})
+        self.assertEqual(len(result["job-1"]), 1)
+        self.assertEqual(len(result["job-2"]), 1)
+
+    async def test_preserves_rgba_alpha(self):
+        output = {"images": [self.artifact("render_00001_.png", png_bytes(mode="RGBA"))]}
+
+        result = await media.collect_images(output, self.manifest())
+
+        self.assertEqual(result["job-1"][0].shape, (1, 2, 4))
+        self.assertAlmostEqual(float(result["job-1"][0][0, 0, 3]), 40 / 255)
+
+    async def test_signed_url_is_downloaded_only_through_injected_downloader(self):
+        output = {"images": [{
+            "filename": "render_00001_.png", "artifact_type": "output", "type": "gcs_url",
+            "url": "https://storage.example.test/signed",
+        }]}
+        fetch = AsyncMock(return_value=png_bytes())
+
+        result = await media.collect_images(output, self.manifest(), download=fetch)
+
+        self.assertEqual(result["job-1"][0].shape, (1, 2, 3))
+        fetch.assert_awaited_once_with("https://storage.example.test/signed")
+
+    async def test_ignores_unrelated_artifacts_but_requires_every_manifest_collector(self):
+        output = {"images": [
+            self.artifact("preview_00001_.png"),
+            self.artifact("render_00001_.png"),
+            {"filename": "mesh.obj", "type": "url", "url": "not a fetch target"},
+        ]}
+
+        result = await media.collect_images(output, self.manifest())
+
+        self.assertEqual(len(result["job-1"]), 1)
+        with self.assertRaisesRegex(media.ServerlessMediaError, "expected|missing"):
+            await media.collect_images({"images": [self.artifact("preview_00001_.png")]}, self.manifest())
+
+    async def test_accepts_compiler_length_flat_safe_prefix(self):
+        prefix = "serverless_" + "j" * 24 + "_" + "w" * 24 + "_" + "c" * 24 + "_image_" + "a" * 64
+        manifest = self.manifest(prefix)
+        output = {"images": [self.artifact(f"{prefix}_00001_.png")]}
+
+        result = await media.collect_images(output, manifest)
+
+        self.assertEqual(len(result["job-1"]), 1)
+
+    async def test_rejects_duplicate_manifest_prefix_and_unsafe_prefix(self):
+        duplicate = self.manifest() * 2
+        with self.assertRaises(media.ServerlessMediaError):
+            await media.collect_images({"images": []}, duplicate)
+        unsafe = self.manifest("../render")
+        with self.assertRaises(media.ServerlessMediaError):
+            await media.collect_images({"images": []}, unsafe)
+
+    async def test_rejects_duplicate_canonical_artifact_identity(self):
+        item = self.artifact("render_00001_.png")
+        with self.assertRaisesRegex(media.ServerlessMediaError, "[Dd]uplicate"):
+            await media.collect_images({"images": [item, dict(item)]}, self.manifest())
+
+    async def test_rejects_matching_non_output_artifact_and_non_numeric_counter(self):
+        wrong_type = self.artifact("render_00001_.png", artifact_type="temp")
+        with self.assertRaises(media.ServerlessMediaError):
+            await media.collect_images({"images": [wrong_type]}, self.manifest())
+        with self.assertRaises(media.ServerlessMediaError):
+            await media.collect_images({"images": [self.artifact("render_final.png")]}, self.manifest())
+
+    async def test_rejects_provider_errors_even_when_job_completed(self):
+        for output in (
+            {"error": "raw provider secret", "images": []},
+            {"errors": ["raw provider secret"], "images": []},
+            {"status": "COMPLETED", "output": {"error": "raw provider secret", "images": []}},
+        ):
+            with self.subTest(output=output):
+                with self.assertRaises(media.ServerlessMediaError) as ctx:
+                    await media.collect_images(output, self.manifest())
+                self.assertNotIn("raw provider secret", str(ctx.exception))
+
+    async def test_rejects_invalid_base64_and_malformed_image(self):
+        for item in (
+            self.artifact("render_00001_.png", data="!not base64!"),
+            self.artifact("render_00001_.png", data=base64.b64encode(b"not an image").decode()),
+        ):
+            with self.subTest(item=item):
+                with self.assertRaises(media.ServerlessMediaError):
+                    await media.collect_images({"images": [item]}, self.manifest())
+
+    async def test_validates_shape_consistency_within_collector(self):
+        output = {"images": [
+            self.artifact("render_00001_.png", png_bytes(size=(2, 1))),
+            self.artifact("render_00002_.png", png_bytes(size=(3, 1))),
+        ]}
+        with self.assertRaisesRegex(media.ServerlessMediaError, "shape|dimensions"):
+            await media.collect_images(output, self.manifest())
+
+    async def test_bounds_output_file_count_and_bytes(self):
+        too_many = {"images": [{}] * (media.MAX_OUTPUT_FILES + 1)}
+        with self.assertRaisesRegex(media.ServerlessMediaError, "files"):
+            await media.collect_images(too_many, self.manifest())
+
+        with patch.object(media, "MAX_OUTPUT_IMAGE_BYTES", 4):
+            output = {"images": [self.artifact("render_00001_.png", data=base64.b64encode(b"12345").decode())]}
+            with self.assertRaisesRegex(media.ServerlessMediaError, "size"):
+                await media.collect_images(output, self.manifest())
+
+    async def test_aggregate_payload_limit_is_checked_as_artifacts_are_read(self):
+        output = {"images": [
+            self.artifact("render_00001_.png", data=base64.b64encode(b"12").decode()),
+            self.artifact("render_00002_.png", data=base64.b64encode(b"34").decode()),
+        ]}
+        with patch.object(media, "MAX_OUTPUT_TOTAL_BYTES", 3):
+            with self.assertRaisesRegex(media.ServerlessMediaError, "size"):
+                await media.collect_images(output, self.manifest())
+
+    async def test_rejects_non_public_or_non_https_artifact_urls_before_downloading(self):
+        for url in ("http://example.test/a", "https://127.0.0.1/a", "https://user:pass@example.test/a"):
+            with self.subTest(url=url):
+                fetch = AsyncMock(return_value=png_bytes())
+                output = {"images": [{"filename": "render_00001_.png", "type": "url", "url": url}]}
+                with self.assertRaises(media.ServerlessMediaError):
+                    await media.collect_images(output, self.manifest(), download=fetch)
+                fetch.assert_not_awaited()
+
+    async def test_public_resolver_rejects_non_global_dns_answers(self):
+        resolver = media._PublicResolver()
+        resolver._getaddrinfo = AsyncMock(return_value=[(2, 1, 6, "", ("8.8.8.8", 443))])
+        self.assertEqual(len(await resolver.resolve("cdn.example.test", 443)), 1)
+        resolver._getaddrinfo = AsyncMock(return_value=[(2, 1, 6, "", ("10.0.0.5", 443))])
+        with self.assertRaises(OSError):
+            await resolver.resolve("cdn.example.test", 443)
+
+    async def test_default_download_revalidates_redirects_and_disables_environment_proxies(self):
+        class Content:
+            def __init__(self, chunks):
+                self.chunks = iter(chunks)
+            async def iter_chunked(self, _size):
+                for chunk in self.chunks:
+                    yield chunk
+
+        class Response:
+            def __init__(self, status, headers=None, chunks=()):
+                self.status, self.headers = status, headers or {}
+                self.content = Content(chunks)
+            async def __aenter__(self):
+                return self
+            async def __aexit__(self, *_args):
+                return False
+
+        class Session:
+            def __init__(self, responses):
+                self.responses, self.calls = list(responses), []
+            async def __aenter__(self):
+                return self
+            async def __aexit__(self, *_args):
+                return False
+            def get(self, url, **kwargs):
+                self.calls.append((url, kwargs))
+                return self.responses.pop(0)
+
+        session = Session([Response(302, {"Location": "https://cdn.example.test/final"}), Response(200, chunks=[b"abc", b"def"])])
+        with patch.object(media.aiohttp, "ClientSession", return_value=session) as create_session:
+            payload = await media.download_public_artifact("https://storage.example.test/signed")
+
+        self.assertEqual(payload, b"abcdef")
+        self.assertEqual([call[0] for call in session.calls], ["https://storage.example.test/signed", "https://cdn.example.test/final"])
+        self.assertTrue(all(call[1]["allow_redirects"] is False for call in session.calls))
+        self.assertFalse(create_session.call_args.kwargs["trust_env"])
+        self.assertIsInstance(create_session.call_args.kwargs["connector"]._resolver, media._PublicResolver)
+        self.assertNotIn("Authorization", session.calls[0][1].get("headers", {}))
+        await create_session.call_args.kwargs["connector"].close()
+
+        blocked = Session([Response(302, {"Location": "https://127.0.0.1/metadata"})])
+        with patch.object(media.aiohttp, "ClientSession", return_value=blocked) as blocked_session:
+            with self.assertRaises(media.ServerlessMediaError):
+                await media.download_public_artifact("https://storage.example.test/signed")
+        self.assertEqual(len(blocked.calls), 1)
+        await blocked_session.call_args.kwargs["connector"].close()
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/test_serverless_prompt.py

@@ -0,0 +1,392 @@
+import importlib.util
+import json
+import sys
+import types
+from pathlib import Path
+
+import pytest
+
+
+MODULE_PATH = (
+    Path(__file__).resolve().parents[1]
+    / "api"
+    / "orchestration"
+    / "serverless_prompt.py"
+)
+_spec = importlib.util.spec_from_file_location("serverless_prompt_under_test", MODULE_PATH)
+assert _spec is not None and _spec.loader is not None
+_serverless_prompt = importlib.util.module_from_spec(_spec)
+_spec.loader.exec_module(_serverless_prompt)
+compile_serverless_prompt = _serverless_prompt.compile_serverless_prompt
+
+
+def _prepared_prompt():
+    return {
+        "1": {"class_type": "CheckpointLoaderSimple", "inputs": {}, "_meta": {"title": "source"}},
+        "2": {"class_type": "KSampler", "inputs": {"model": ["1", 0], "seed": ["3", 0]}},
+        "3": {
+            "class_type": "DistributedSeed",
+            "inputs": {"seed": 40, "is_worker": True, "worker_id": "worker_0"},
+        },
+        "4": {
+            "class_type": "DistributedCollector",
+            "inputs": {"images": ["2", 0], "multi_job_id": "job/42", "worker_id": "worker_0", "is_worker": True},
+        },
+        "5": {"class_type": "VAEDecode", "inputs": {"samples": ["4", 0]}},
+        "6": {"class_type": "SaveImage", "inputs": {"images": ["5", 0]}},
+        "7": {"class_type": "PreviewImage", "inputs": {"images": ["4", 0]}, "_meta": {"title": "Preview Image (auto-added)"}},
+        "8": {"class_type": "UnrelatedNode", "inputs": {}},
+    }
+
+
+def _run(prompt):
+    return compile_serverless_prompt(prompt)
+
+
+def _load_prompt_transform():
+    module_path = MODULE_PATH.with_name("prompt_transform.py")
+    package_name = "serverless_prompt_transform_testpkg"
+    root_pkg = types.ModuleType(package_name)
+    root_pkg.__path__ = []
+    sys.modules[package_name] = root_pkg
+    for suffix in ("api", "api.orchestration", "utils"):
+        pkg = types.ModuleType(f"{package_name}.{suffix}")
+        pkg.__path__ = []
+        sys.modules[f"{package_name}.{suffix}"] = pkg
+    logging_module = types.ModuleType(f"{package_name}.utils.logging")
+    setattr(logging_module, "debug_log", lambda *_args, **_kwargs: None)
+    setattr(logging_module, "log", lambda *_args, **_kwargs: None)
+    sys.modules[f"{package_name}.utils.logging"] = logging_module
+    spec = importlib.util.spec_from_file_location(
+        f"{package_name}.api.orchestration.prompt_transform", module_path
+    )
+    assert spec is not None and spec.loader is not None
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    return module
+
+
+def test_compiles_prompt_prepared_by_existing_worker_helpers():
+    prompt_transform = _load_prompt_transform()
+    original = {
+        "1": {"class_type": "CheckpointLoaderSimple", "inputs": {}},
+        "2": {"class_type": "KSampler", "inputs": {"model": ["1", 0], "seed": ["3", 0]}},
+        "3": {"class_type": "DistributedSeed", "inputs": {"seed": 100}},
+        "4": {"class_type": "DistributedCollector", "inputs": {"images": ["2", 0]}},
+        "5": {"class_type": "SaveImage", "inputs": {"images": ["4", 0]}},
+    }
+    pruned = prompt_transform.prune_prompt_for_worker(original)
+    prompt_index = prompt_transform.PromptIndex(pruned)
+    prepared = prompt_transform.apply_participant_overrides(
+        pruned,
+        participant_id="worker-a",
+        enabled_worker_ids=["worker-a"],
+        job_id_map=prompt_transform.generate_job_id_map(prompt_index, "run"),
+        master_url="http://master.invalid",
+        delegate_master=False,
+        prompt_index=prompt_index,
+    )
+
+    graph, manifest = _run(prepared)
+
+    assert graph["2"]["inputs"]["seed"] == 101
+    assert set(graph) == {"1", "2", "7"}
+    assert graph["7"]["class_type"] == "SaveImage"
+    assert manifest[0]["job_id"] == "run_4"
+    assert manifest[0]["collector_id"] == "4"
+
+
+def test_compiles_collector_input_to_flat_saveimage_and_prunes_downstream():
+    original = _prepared_prompt()
+
+    graph, manifest = _run(original)
+
+    assert manifest == [
+        {
+            "collector_id": "4",
+            "job_id": "job/42",
+            "node_id": "9",
+            "kind": "image",
+            "prefix": graph["9"]["inputs"]["filename_prefix"],
+        }
+    ]
+    prefix = manifest[0]["prefix"]
+    assert graph["9"] == {
+        "class_type": "SaveImage",
+        "inputs": {"images": ["2", 0], "filename_prefix": prefix},
+        "_meta": {"title": "Serverless output: collector 4 (image)"},
+    }
+    assert "/" not in prefix and "\\" not in prefix
+    assert graph["2"]["inputs"]["seed"] == 41
+    assert set(graph) == {"1", "2", "9"}
+    assert all("Distributed" not in node["class_type"] for node in graph.values())
+    assert original["3"]["inputs"]["seed"] == 40
+    assert original["4"]["inputs"]["images"] == ["2", 0]
+
+
+def test_manifest_and_prefix_are_repeatable_and_safe_for_untrusted_identity_parts():
+    prompt = _prepared_prompt()
+    prompt["4"]["inputs"].update(multi_job_id="../run/a", worker_id="worker/unsafe")
+
+    first_graph, first_manifest = _run(prompt)
+    second_graph, second_manifest = _run(prompt)
+
+    assert first_manifest == second_manifest
+    assert first_graph == second_graph
+    prefix = first_manifest[0]["prefix"]
+    assert prefix == first_graph[first_manifest[0]["node_id"]]["inputs"]["filename_prefix"]
+    assert not any(char in prefix for char in ("/", "\\", ":"))
+    assert first_manifest[0]["job_id"] == "../run/a"
+
+
+def test_compiles_multiple_disconnected_collectors_in_stable_order():
+    prompt = {
+        "10": {"class_type": "LoadImage", "inputs": {}},
+        "20": {"class_type": "LoadImage", "inputs": {}},
+        "2": {"class_type": "DistributedCollector", "inputs": {"images": ["20", 0], "multi_job_id": "b", "worker_id": "worker_1"}},
+        "1": {"class_type": "DistributedCollector", "inputs": {"images": ["10", 0], "multi_job_id": "a", "worker_id": "worker_0"}},
+        "30": {"class_type": "PreviewImage", "inputs": {"images": ["1", 0]}},
+    }
+
+    graph, manifest = _run(prompt)
+
+    assert [item["collector_id"] for item in manifest] == ["1", "2"]
+    assert [item["node_id"] for item in manifest] == ["31", "32"]
+    assert graph["31"]["inputs"]["images"] == ["10", 0]
+    assert graph["32"]["inputs"]["images"] == ["20", 0]
+    assert set(graph) == {"10", "20", "31", "32"}
+
+
+def test_rejects_audio_input_instead_of_emitting_audio_saver():
+    prompt = {
+        "1": {"class_type": "LoadAudio", "inputs": {}},
+        "2": {"class_type": "DistributedCollector", "inputs": {"audio": ["1", 0]}},
+    }
+
+    with pytest.raises(ValueError, match="audio.*unsupported"):
+        _run(prompt)
+
+
+def test_rejects_collector_without_image_input():
+    with pytest.raises(ValueError, match="collector.*image"):
+        _run({"1": {"class_type": "DistributedCollector", "inputs": {}}})
+
+
+@pytest.mark.parametrize(
+    "class_type",
+    ["UltimateSDUpscaleDistributed", "DistributedEmptyImage", "ImageBatchDivider", "AudioBatchDivider"],
+)
+def test_rejects_unsupported_distributed_and_batch_nodes(class_type):
+    prompt = {
+        "1": {"class_type": class_type, "inputs": {}},
+        "2": {"class_type": "DistributedCollector", "inputs": {"images": ["3", 0]}},
+        "3": {"class_type": "LoadImage", "inputs": {}},
+    }
+
+    with pytest.raises(ValueError, match="unsupported"):
+        _run(prompt)
+
+
+def test_rejects_collector_to_collector_dependency_through_processing_node():
+    prompt = {
+        "1": {"class_type": "LoadImage", "inputs": {}},
+        "2": {"class_type": "DistributedCollector", "inputs": {"images": ["1", 0]}},
+        "3": {"class_type": "ImageScale", "inputs": {"image": ["2", 0]}},
+        "4": {"class_type": "DistributedCollector", "inputs": {"images": ["3", 0]}},
+    }
+
+    with pytest.raises(ValueError, match="collector-to-collector"):
+        _run(prompt)
+
+
+@pytest.mark.parametrize(
+    "link",
+    [["missing", 0], ["1", -1], ["1", "0"], ["1", True]],
+)
+def test_rejects_malformed_api_links(link):
+    prompt = {
+        "1": {"class_type": "LoadImage", "inputs": {}},
+        "2": {"class_type": "DistributedCollector", "inputs": {"images": link}},
+    }
+
+    with pytest.raises(ValueError, match="link"):
+        _run(prompt)
+
+
+def test_rejects_unexpected_output_slot_from_removed_distributed_node():
+    prompt = {
+        "1": {"class_type": "DistributedSeed", "inputs": {"seed": 1}},
+        "2": {"class_type": "KSampler", "inputs": {"seed": ["1", 1]}},
+        "3": {"class_type": "DistributedCollector", "inputs": {"images": ["2", 0]}},
+    }
+
+    with pytest.raises(ValueError, match="output slot"):
+        _run(prompt)
+
+
+def test_distributed_seed_matches_worker_offset_and_master_passthrough():
+    prompt = {
+        "1": {"class_type": "DistributedSeed", "inputs": {"seed": 10, "is_worker": True, "worker_id": "worker_1"}},
+        "2": {"class_type": "DistributedSeed", "inputs": {"seed": 7, "is_worker": False}},
+        "3": {"class_type": "KSampler", "inputs": {"seed": ["1", 0], "steps": ["2", 0]}},
+        "4": {"class_type": "DistributedCollector", "inputs": {"images": ["3", 0]}},
+    }
+
+    graph, _manifest = _run(prompt)
+
+    assert graph["3"]["inputs"] == {"seed": 12, "steps": 7}
+    assert "1" not in graph and "2" not in graph
+
+
+def test_distributed_value_uses_one_based_worker_values_and_type_coercion():
+    prompt = {
+        "1": {"class_type": "DistributedValue", "inputs": {"default_value": "3.9", "worker_values": json.dumps({"_type": "INT", "1": "8.7"}), "is_worker": True, "worker_id": "worker_0"}},
+        "2": {"class_type": "DistributedValue", "inputs": {"default_value": "3.9", "worker_values": json.dumps({"_type": "INT", "2": ""}), "is_worker": True, "worker_id": "worker_1"}},
+        "3": {"class_type": "KSampler", "inputs": {"steps": ["1", 0], "seed": ["2", 0]}},
+        "4": {"class_type": "DistributedCollector", "inputs": {"images": ["3", 0]}},
+    }
+
+    graph, _manifest = _run(prompt)
+
+    assert graph["3"]["inputs"] == {"steps": 8, "seed": 3}
+    assert "1" not in graph and "2" not in graph
+
+
+def test_resolves_stock_primitive_links_for_distributed_seed_and_value():
+    prompt = {
+        "1": {"class_type": "PrimitiveInt", "inputs": {"value": 25}},
+        "2": {"class_type": "DistributedSeed", "inputs": {"seed": ["1", 0], "is_worker": True, "worker_id": "worker_0"}},
+        "3": {"class_type": "PrimitiveString", "inputs": {"value": "4.8"}},
+        "4": {"class_type": "PrimitiveString", "inputs": {"value": '{"_type":"FLOAT","1":"2.75"}'}},
+        "5": {"class_type": "DistributedValue", "inputs": {"default_value": ["3", 0], "worker_values": ["4", 0], "is_worker": True, "worker_id": "worker_0"}},
+        "6": {"class_type": "KSampler", "inputs": {"seed": ["2", 0], "cfg": ["5", 0]}},
+        "7": {"class_type": "DistributedCollector", "inputs": {"images": ["6", 0]}},
+    }
+
+    graph, _manifest = _run(prompt)
+
+    assert graph["6"]["inputs"] == {"seed": 26, "cfg": 2.75}
+    assert not {"1", "2", "3", "4", "5"}.intersection(graph)
+
+
+def test_rejects_dynamic_linked_seed_or_value_without_executing_it():
+    prompt = {
+        "1": {"class_type": "ArbitraryNode", "inputs": {}},
+        "2": {"class_type": "DistributedSeed", "inputs": {"seed": ["1", 0], "is_worker": True, "worker_id": "worker_0"}},
+        "3": {"class_type": "KSampler", "inputs": {"seed": ["2", 0]}},
+        "4": {"class_type": "DistributedCollector", "inputs": {"images": ["3", 0]}},
+    }
+
+    with pytest.raises(ValueError, match="literal or supported Primitive"):
+        _run(prompt)
+
+
+def test_primitive_resolution_detects_cycles():
+    prompt = {
+        "1": {"class_type": "PrimitiveInt", "inputs": {"value": ["2", 0]}},
+        "2": {"class_type": "PrimitiveInt", "inputs": {"value": ["1", 0]}},
+        "3": {"class_type": "DistributedSeed", "inputs": {"seed": ["1", 0]}},
+        "4": {"class_type": "KSampler", "inputs": {"seed": ["3", 0]}},
+        "5": {"class_type": "DistributedCollector", "inputs": {"images": ["4", 0]}},
+    }
+
+    with pytest.raises(ValueError, match="cycle"):
+        _run(prompt)
+
+
+def test_model_name_trivial_text_passthrough_is_inlined():
+    prompt = {
+        "1": {"class_type": "DistributedModelName", "inputs": {"text": "model.safetensors"}},
+        "2": {"class_type": "KSampler", "inputs": {"model_name": ["1", 0]}},
+        "3": {"class_type": "DistributedCollector", "inputs": {"images": ["2", 0]}},
+    }
+
+    graph, _manifest = _run(prompt)
+
+    assert graph["2"]["inputs"]["model_name"] == "model.safetensors"
+    assert "1" not in graph
+
+
+def test_rejects_dynamic_linked_model_name_text():
+    prompt = {
+        "1": {"class_type": "ArbitraryNode", "inputs": {}},
+        "2": {"class_type": "DistributedModelName", "inputs": {"text": ["1", 0]}},
+        "3": {"class_type": "KSampler", "inputs": {"model_name": ["2", 0]}},
+        "4": {"class_type": "DistributedCollector", "inputs": {"images": ["3", 0]}},
+    }
+
+    with pytest.raises(ValueError, match="DistributedModelName.*literal or supported Primitive"):
+        _run(prompt)
+
+
+def test_resolves_boolean_and_float_primitive_links():
+    prompt = {
+        "1": {"class_type": "PrimitiveBoolean", "inputs": {"value": True}},
+        "2": {"class_type": "PrimitiveFloat", "inputs": {"value": 1.25}},
+        "3": {"class_type": "DistributedValue", "inputs": {"default_value": ["1", 0], "is_worker": False}},
+        "4": {"class_type": "DistributedValue", "inputs": {"default_value": ["2", 0], "worker_values": '{"_type":"FLOAT","1":"2.5"}', "is_worker": True, "worker_id": "worker_0"}},
+        "5": {"class_type": "KSampler", "inputs": {"enabled": ["3", 0], "cfg": ["4", 0]}},
+        "6": {"class_type": "DistributedCollector", "inputs": {"images": ["5", 0]}},
+    }
+
+    graph, _manifest = _run(prompt)
+
+    assert graph["5"]["inputs"] == {"enabled": True, "cfg": 2.5}
+    assert not {"1", "2", "3", "4"}.intersection(graph)
+
+
+def test_compiled_graph_matches_golden_prompt_and_manifest():
+    prompt = _prepared_prompt()
+    prompt["3"]["inputs"]["seed"] = ["10", 0]
+    prompt["10"] = {"class_type": "PrimitiveInt", "inputs": {"value": 40}}
+    prompt["11"] = {"class_type": "PrimitiveFloat", "inputs": {"value": 1.25}}
+    prompt["12"] = {
+        "class_type": "DistributedValue",
+        "inputs": {
+            "default_value": ["11", 0],
+            "worker_values": json.dumps({"_type": "FLOAT", "1": "2.5"}),
+            "is_worker": True,
+            "worker_id": "worker_0",
+        },
+    }
+    prompt["13"] = {"class_type": "PrimitiveString", "inputs": {"value": "model.safetensors"}}
+    prompt["14"] = {"class_type": "DistributedModelName", "inputs": {"text": ["13", 0]}}
+    prompt["2"]["inputs"].update(cfg=["12", 0], model_name=["14", 0])
+
+    graph, manifest = _run(prompt)
+
+    expected_prefix = (
+        "serverless_job_42_worker_0_4_image_"
+        "1c81da108a2eb60282bd436512acfc59903aba73e49b19b6f47bbdcafd30d244"
+    )
+    assert list(graph) == ["1", "2", "15"]
+    assert graph == {
+        "1": {
+            "class_type": "CheckpointLoaderSimple",
+            "inputs": {},
+            "_meta": {"title": "source"},
+        },
+        "2": {
+            "class_type": "KSampler",
+            "inputs": {
+                "model": ["1", 0],
+                "seed": 41,
+                "cfg": 2.5,
+                "model_name": "model.safetensors",
+            },
+        },
+        "15": {
+            "class_type": "SaveImage",
+            "inputs": {"images": ["2", 0], "filename_prefix": expected_prefix},
+            "_meta": {"title": "Serverless output: collector 4 (image)"},
+        },
+    }
+    assert manifest == [
+        {
+            "collector_id": "4",
+            "job_id": "job/42",
+            "node_id": "15",
+            "kind": "image",
+            "prefix": expected_prefix,
+        }
+    ]

tests/test_serverless_routes.py

@@ -0,0 +1,624 @@
+import copy
+import importlib.util
+import json
+import sys
+import tempfile
+import types
+import unittest
+from pathlib import Path
+from unittest.mock import patch
+
+
+class _FakeResponse:
+    def __init__(self, payload, status=200):
+        self.payload = copy.deepcopy(payload)
+        self.status = status
+
+
+class _FakeRequest:
+    def __init__(self, payload=None, error=None):
+        self.payload = payload
+        self.error = error
+
+    async def json(self):
+        if self.error is not None:
+            raise self.error
+        return self.payload
+
+
+class _Routes:
+    def get(self, _path):
+        return lambda handler: handler
+
+    def post(self, _path):
+        return lambda handler: handler
+
+
+def _load_modules():
+    repo = Path(__file__).resolve().parents[1]
+    package_name = "dist_serverless_routes_testpkg"
+    for name in list(sys.modules):
+        if name == package_name or name.startswith(f"{package_name}."):
+            del sys.modules[name]
+
+    root = types.ModuleType(package_name)
+    root.__path__ = [str(repo)]
+    sys.modules[package_name] = root
+    api_pkg = types.ModuleType(f"{package_name}.api")
+    api_pkg.__path__ = [str(repo / "api")]
+    sys.modules[f"{package_name}.api"] = api_pkg
+    utils_pkg = types.ModuleType(f"{package_name}.utils")
+    utils_pkg.__path__ = [str(repo / "utils")]
+    sys.modules[f"{package_name}.utils"] = utils_pkg
+    orchestration_pkg = types.ModuleType(f"{package_name}.api.orchestration")
+    orchestration_pkg.__path__ = [str(repo / "api" / "orchestration")]
+    sys.modules[f"{package_name}.api.orchestration"] = orchestration_pkg
+
+    logging_module = types.ModuleType(f"{package_name}.utils.logging")
+    logging_module.log = lambda *_args, **_kwargs: None
+    logging_module.debug_log = lambda *_args, **_kwargs: None
+    sys.modules[logging_module.__name__] = logging_module
+    constants_module = types.ModuleType(f"{package_name}.utils.constants")
+    constants_module.HEARTBEAT_TIMEOUT = 30
+    sys.modules[constants_module.__name__] = constants_module
+
+    def load_file(module_name, path):
+        spec = importlib.util.spec_from_file_location(module_name, path)
+        assert spec is not None and spec.loader is not None
+        module = importlib.util.module_from_spec(spec)
+        sys.modules[module_name] = module
+        spec.loader.exec_module(module)
+        return module
+
+    config_module = load_file(f"{package_name}.utils.config", repo / "utils" / "config.py")
+    serverless_config = load_file(
+        f"{package_name}.utils.serverless_config", repo / "utils" / "serverless_config.py"
+    )
+
+    runtime = types.ModuleType(f"{package_name}.api.orchestration.serverless_runtime")
+    runtime.secret_store = None
+    runtime.manager = types.SimpleNamespace(tasks={}, runs={}, pool_is_active=lambda _: False)
+    runtime.get_secret_store = lambda: runtime.secret_store
+    runtime.get_manager = lambda: runtime.manager
+    sys.modules[runtime.__name__] = runtime
+
+    import aiohttp  # Ensure the real dependency is loaded; never replace global aiohttp.
+    previous_server = sys.modules.get("server")
+    server = types.ModuleType("server")
+    server.PromptServer = types.SimpleNamespace(instance=types.SimpleNamespace(routes=_Routes()))
+    sys.modules["server"] = server
+
+    network = types.ModuleType(f"{package_name}.utils.network")
+
+    async def handle_api_error(_request, error, status=500):
+        return _FakeResponse({"status": "error", "message": str(error)}, status)
+
+    network.handle_api_error = handle_api_error
+    network.normalize_host = lambda value: value
+    sys.modules[network.__name__] = network
+
+    routes = load_file(f"{package_name}.api.config_routes", repo / "api" / "config_routes.py")
+    routes.web = types.SimpleNamespace(
+        json_response=lambda payload, status=200: _FakeResponse(payload, status)
+    )
+
+    if previous_server is None:
+        sys.modules.pop("server", None)
+    else:
+        sys.modules["server"] = previous_server
+    return config_module, serverless_config, runtime, routes
+
+
+config_module, serverless_config, runtime_module, config_routes = _load_modules()
+
+
+class _RunningTask:
+    def done(self):
+        return False
+
+
+class _ConfigRouteFixture(unittest.IsolatedAsyncioTestCase):
+    def setUp(self):
+        self.tempdir = tempfile.TemporaryDirectory()
+        self.root = Path(self.tempdir.name)
+        config_module.CONFIG_FILE = str(self.root / "gpu_config.json")
+        config_module.invalidate_config_cache()
+        self.config = config_module.get_default_config()
+        self.config["workers"] = []
+        self.assertTrue(config_module.save_config(self.config))
+        self.secrets = serverless_config.SecretStore(self.root / "private" / "credentials.json")
+        runtime_module.secret_store = self.secrets
+        runtime_module.manager = types.SimpleNamespace(tasks={}, runs={}, pool_is_active=lambda _: False)
+
+    def tearDown(self):
+        config_module.invalidate_config_cache()
+        self.tempdir.cleanup()
+
+    def read_config_file(self):
+        return json.loads(Path(config_module.CONFIG_FILE).read_text(encoding="utf-8"))
+
+
+class ServerlessConfigRouteTests(_ConfigRouteFixture):
+    async def test_get_config_redacts_keys_and_reports_only_configured_boolean(self):
+        self.secrets.set("pool-1", "private-api-key")
+        self.config["workers"] = [
+            {"id": "pool-1", "type": "serverless", "provider": "runpod", "endpoint_id": "ep-1",
+             "api_key": "inline-api-key", "name": "GPU"}
+        ]
+        Path(config_module.CONFIG_FILE).write_text(json.dumps(self.config), encoding="utf-8")
+        config_module.invalidate_config_cache()
+
+        response = await config_routes.get_config_endpoint(_FakeRequest())
+
+        self.assertEqual(response.status, 200)
+        worker = response.payload["workers"][0]
+        self.assertTrue(worker["api_key_configured"])
+        self.assertNotIn("api_key", worker)
+        self.assertNotIn("private-api-key", json.dumps(response.payload))
+        self.assertNotIn("inline-api-key", json.dumps(response.payload))
+
+    async def test_get_config_redacts_normalized_secret_field_aliases(self):
+        aliases = {
+            "apiKey": "camel-api-key",
+            "APIKEY": "upper-api-key",
+            "clearApiKey": "camel-clear-key",
+            "CLEAR_API_KEY": "upper-clear-key",
+            "runpod_CREDENTIAL": "provider-credential",
+        }
+        self.config["workers"] = [{
+            "id": "pool-1", "type": "serverless", "endpoint_id": "ep-123",
+            **aliases,
+            "nested": {"credential_note": "nested-credential"},
+        }]
+        Path(config_module.CONFIG_FILE).write_text(json.dumps(self.config), encoding="utf-8")
+        config_module.invalidate_config_cache()
+
+        response = await config_routes.get_config_endpoint(_FakeRequest())
+
+        self.assertEqual(response.status, 200)
+        worker = response.payload["workers"][0]
+        for alias, value in aliases.items():
+            self.assertNotIn(alias, worker)
+            self.assertNotIn(value, json.dumps(response.payload))
+        self.assertNotIn("credential_note", worker["nested"])
+        self.assertNotIn("nested-credential", json.dumps(response.payload))
+
+    async def test_bulk_config_rejects_normalized_secret_field_aliases(self):
+        for secret_field in ("apiKey", "APIKEY", "clearApiKey", "RUNPOD_CREDENTIAL"):
+            with self.subTest(secret_field=secret_field):
+                secret = f"private-{secret_field}"
+                response = await config_routes.update_config_endpoint(_FakeRequest({
+                    "workers": [{"id": "pool-1", "type": "local", secret_field: secret}]
+                }))
+                self.assertEqual(response.status, 400)
+                self.assertNotIn(secret, json.dumps(response.payload))
+        self.assertEqual(self.read_config_file()["workers"], [])
+
+    async def test_create_serverless_worker_without_local_fields_and_store_key_privately(self):
+        secret = "runpod-secret-never-returned"
+        payload = {
+            "worker_id": "pool-1", "name": "GPU pool", "type": "serverless", "provider": "runpod",
+            "endpoint_id": "ep-123", "jobs_per_run": 3, "concurrency": 2,
+            "timeout_seconds": 600, "enabled": True, "api_key": secret,
+        }
+
+        response = await config_routes.update_worker_endpoint(_FakeRequest(payload))
+
+        self.assertEqual(response.status, 200)
+        self.assertTrue(response.payload["api_key_configured"])
+        self.assertNotIn(secret, json.dumps(response.payload))
+        saved = self.read_config_file()["workers"][0]
+        self.assertEqual(saved["endpoint_id"], "ep-123")
+        self.assertEqual(saved["jobs_per_run"], 3)
+        # Legacy field from older browser tabs is accepted but not persisted.
+        self.assertNotIn("concurrency", saved)
+        self.assertEqual(saved["timeout_seconds"], 600)
+        for obsolete in ("host", "port", "cuda_device", "extra_args", "api_key", "clear_api_key"):
+            self.assertNotIn(obsolete, saved)
+        self.assertEqual(self.secrets.get("pool-1"), secret)
+
+    async def test_enable_only_update_preserves_serverless_settings_and_saved_key(self):
+        self.secrets.set("pool-1", "saved-secret")
+        self.config["workers"] = [{
+            "id": "pool-1", "type": "serverless", "provider": "runpod", "name": "Pool",
+            "endpoint_id": "ep-123", "jobs_per_run": 4, "concurrency": 3,
+            "timeout_seconds": 900, "enabled": False,
+        }]
+        self.assertTrue(config_module.save_config(self.config))
+
+        response = await config_routes.update_worker_endpoint(
+            _FakeRequest({"worker_id": "pool-1", "enabled": True})
+        )
+
+        worker = self.read_config_file()["workers"][0]
+        self.assertTrue(worker["enabled"])
+        self.assertEqual(worker["endpoint_id"], "ep-123")
+        self.assertEqual(worker["jobs_per_run"], 4)
+        self.assertNotIn("concurrency", worker)
+        self.assertEqual(worker["timeout_seconds"], 900)
+        self.assertEqual(self.secrets.get("pool-1"), "saved-secret")
+        self.assertTrue(response.payload["api_key_configured"])
+
+    async def test_invalid_serverless_record_is_not_persisted_or_secret_written(self):
+        secret = "must-not-leak"
+        response = await config_routes.update_worker_endpoint(_FakeRequest({
+            "worker_id": "pool-1", "type": "serverless", "provider": "runpod",
+            "endpoint_id": "ep-123", "jobs_per_run": True, "api_key": secret,
+        }))
+
+        self.assertEqual(response.status, 400)
+        self.assertNotIn(secret, json.dumps(response.payload))
+        self.assertEqual(self.read_config_file()["workers"], [])
+        self.assertIsNone(self.secrets.get("pool-1"))
+
+    async def test_json_decode_error_is_generic_and_does_not_echo_body(self):
+        secret = "body-secret-value"
+        response = await config_routes.update_worker_endpoint(
+            _FakeRequest(error=ValueError(f"bad json near {secret}"))
+        )
+        self.assertEqual(response.status, 400)
+        self.assertNotIn(secret, json.dumps(response.payload))
+
+    async def test_bulk_config_rejects_api_keys_at_any_depth_without_persisting(self):
+        secret = "bulk-secret-never-write"
+        response = await config_routes.update_config_endpoint(_FakeRequest({
+            "workers": [{
+                "id": "pool-1", "type": "serverless", "provider": "runpod",
+                "endpoint_id": "ep-123", "nested": {"api_key": secret},
+            }]
+        }))
+
+        self.assertEqual(response.status, 400)
+        self.assertNotIn(secret, json.dumps(response.payload))
+        self.assertEqual(self.read_config_file()["workers"], [])
+        self.assertFalse(self.secrets.path.exists())
+
+    async def test_bulk_update_validates_and_returns_redacted_serverless_config(self):
+        self.secrets.set("pool-1", "saved-secret")
+        response = await config_routes.update_config_endpoint(_FakeRequest({
+            "workers": [{
+                "id": "pool-1", "type": "serverless", "provider": "runpod",
+                "endpoint_id": "ep-123", "name": "Pool", "enabled": True,
+                "host": "obsolete-host", "port": 8188, "cuda_device": 0,
+            }]
+        }))
+
+        self.assertEqual(response.status, 200)
+        self.assertTrue(response.payload["config"]["workers"][0]["api_key_configured"])
+        saved = self.read_config_file()["workers"][0]
+        for obsolete in ("host", "port", "cuda_device"):
+            self.assertNotIn(obsolete, saved)
+        self.assertNotIn("api_key", response.payload["config"]["workers"][0])
+        self.assertEqual(self.secrets.get("pool-1"), "saved-secret")
+
+    async def test_bulk_type_conversion_clears_secret_and_serverless_fields(self):
+        self.secrets.set("pool-1", "bulk-clear-key")
+        self.config["workers"] = [{
+            "id": "pool-1", "type": "serverless", "provider": "runpod", "endpoint_id": "ep-123",
+        }]
+        self.assertTrue(config_module.save_config(self.config))
+
+        response = await config_routes.update_config_endpoint(_FakeRequest({
+            "workers": [{
+                "id": "pool-1", "type": "local", "port": 8188,
+                "provider": "runpod", "endpoint_id": "legacy-ep", "api_key_configured": True,
+            }]
+        }))
+
+        self.assertEqual(response.status, 200)
+        self.assertIsNone(self.secrets.get("pool-1"))
+        saved = self.read_config_file()["workers"][0]
+        for field in ("provider", "endpoint_id", "api_key_configured"):
+            self.assertNotIn(field, saved)
+
+    async def test_bulk_invalid_serverless_record_does_not_replace_configuration(self):
+        existing = {"id": "local-1", "type": "local", "port": 8188}
+        self.config["workers"] = [existing]
+        self.assertTrue(config_module.save_config(self.config))
+
+        response = await config_routes.update_config_endpoint(_FakeRequest({
+            "workers": [{"id": "pool-1", "type": "serverless", "endpoint_id": "bad/id"}]
+        }))
+
+        self.assertEqual(response.status, 400)
+        self.assertEqual(self.read_config_file()["workers"], [existing])
+
+    async def test_delete_serverless_worker_clears_saved_key(self):
+        self.secrets.set("pool-1", "delete-me")
+        self.config["workers"] = [{
+            "id": "pool-1", "type": "serverless", "provider": "runpod", "endpoint_id": "ep-123"
+        }]
+        self.assertTrue(config_module.save_config(self.config))
+
+        response = await config_routes.delete_worker_endpoint(_FakeRequest({"worker_id": "pool-1"}))
+
+        self.assertEqual(response.status, 200)
+        self.assertEqual(self.read_config_file()["workers"], [])
+        self.assertIsNone(self.secrets.get("pool-1"))
+
+    async def test_type_conversion_clears_secret_and_removes_serverless_fields(self):
+        self.secrets.set("pool-1", "old-key")
+        self.config["workers"] = [{
+            "id": "pool-1", "name": "Pool", "type": "serverless", "provider": "runpod",
+            "endpoint_id": "ep-123", "jobs_per_run": 2, "concurrency": 2, "timeout_seconds": 600,
+        }]
+        self.assertTrue(config_module.save_config(self.config))
+
+        response = await config_routes.update_worker_endpoint(_FakeRequest({
+            "worker_id": "pool-1", "type": "local", "name": "Local", "port": 8188,
+            "cuda_device": 0, "enabled": False,
+        }))
+
+        self.assertEqual(response.status, 200)
+        self.assertIsNone(self.secrets.get("pool-1"))
+        worker = self.read_config_file()["workers"][0]
+        self.assertEqual(worker["type"], "local")
+        for serverless_field in ("provider", "endpoint_id", "jobs_per_run", "concurrency", "timeout_seconds"):
+            self.assertNotIn(serverless_field, worker)
+
+    async def test_enabled_toggle_allowed_during_active_run_but_endpoint_change_is_rejected(self):
+        self.secrets.set("pool-1", "saved-key")
+        self.config["workers"] = [{
+            "id": "pool-1", "type": "serverless", "provider": "runpod", "endpoint_id": "ep-123",
+            "enabled": False,
+        }]
+        self.assertTrue(config_module.save_config(self.config))
+        runtime_module.manager = types.SimpleNamespace(
+            tasks={"run-1": _RunningTask()},
+            runs={"run-1": {"worker": {"id": "serverless:pool-1:1", "pool_id": "pool-1"}}},
+            pool_is_active=lambda worker_id: worker_id == "pool-1",
+        )
+
+        enabled = await config_routes.update_worker_endpoint(
+            _FakeRequest({"worker_id": "pool-1", "enabled": True})
+        )
+        blocked = await config_routes.update_worker_endpoint(_FakeRequest({
+            "worker_id": "pool-1", "endpoint_id": "ep-456",
+        }))
+
+        self.assertEqual(enabled.status, 200)
+        self.assertEqual(blocked.status, 409)
+        self.assertEqual(self.read_config_file()["workers"][0]["endpoint_id"], "ep-123")
+
+    async def test_blank_api_key_keeps_saved_key_and_clear_api_key_deletes_it(self):
+        self.secrets.set("pool-1", "saved-key")
+        self.config["workers"] = [{
+            "id": "pool-1", "type": "serverless", "provider": "runpod", "endpoint_id": "ep-123",
+        }]
+        self.assertTrue(config_module.save_config(self.config))
+
+        keep = await config_routes.update_worker_endpoint(_FakeRequest({
+            "worker_id": "pool-1", "api_key": "",
+        }))
+        clear = await config_routes.update_worker_endpoint(_FakeRequest({
+            "worker_id": "pool-1", "clear_api_key": True,
+        }))
+
+        self.assertTrue(keep.payload["api_key_configured"])
+        self.assertFalse(clear.payload["api_key_configured"])
+        self.assertIsNone(self.secrets.get("pool-1"))
+
+    async def test_config_save_refuses_inline_api_keys_and_invalid_serverless_records(self):
+        with self.subTest(case="inline key"):
+            config = config_module.get_default_config()
+            config["workers"] = [{"id": "pool-1", "type": "local", "api_key": "private"}]
+            self.assertFalse(config_module.save_config(config))
+        with self.subTest(case="invalid serverless record"):
+            config = config_module.get_default_config()
+            config["workers"] = [{"id": "pool-1", "type": "serverless", "endpoint_id": "bad/id"}]
+            self.assertFalse(config_module.save_config(config))
+
+    async def test_active_pool_blocks_credential_changes_type_conversion_and_deletion(self):
+        self.secrets.set("pool-1", "saved-key")
+        self.config["workers"] = [{
+            "id": "pool-1", "type": "serverless", "provider": "runpod", "endpoint_id": "ep-123",
+        }]
+        self.assertTrue(config_module.save_config(self.config))
+        runtime_module.manager = types.SimpleNamespace(
+            tasks={"run-1": _RunningTask()},
+            runs={"run-1": {"worker": {"id": "serverless:pool-1:1", "pool_id": "pool-1"}}},
+            pool_is_active=lambda worker_id: worker_id == "pool-1",
+        )
+
+        credential_change = await config_routes.update_worker_endpoint(_FakeRequest({
+            "worker_id": "pool-1", "api_key": "new-key",
+        }))
+        type_change = await config_routes.update_worker_endpoint(_FakeRequest({
+            "worker_id": "pool-1", "type": "local", "port": 8188, "cuda_device": 0,
+        }))
+        deletion = await config_routes.delete_worker_endpoint(_FakeRequest({"worker_id": "pool-1"}))
+
+        self.assertEqual(credential_change.status, 409)
+        self.assertEqual(type_change.status, 409)
+        self.assertEqual(deletion.status, 409)
+        self.assertEqual(self.secrets.get("pool-1"), "saved-key")
+        self.assertEqual(self.read_config_file()["workers"][0]["type"], "serverless")
+
+    async def test_secret_write_is_rolled_back_when_config_persistence_fails(self):
+        secret = "write-then-rollback"
+        with patch.object(config_module, "save_config", return_value=False):
+            response = await config_routes.update_worker_endpoint(_FakeRequest({
+                "worker_id": "pool-1", "type": "serverless", "endpoint_id": "ep-123",
+                "api_key": secret,
+            }))
+
+        self.assertEqual(response.status, 500)
+        self.assertIsNone(self.secrets.get("pool-1"))
+        self.assertEqual(self.read_config_file()["workers"], [])
+        self.assertEqual(config_module.load_config()["workers"], [])
+
+    async def test_bulk_secret_delete_is_rolled_back_when_config_persistence_fails(self):
+        self.secrets.set("pool-1", "bulk-restore-key")
+        self.config["workers"] = [{
+            "id": "pool-1", "type": "serverless", "endpoint_id": "ep-123",
+        }]
+        self.assertTrue(config_module.save_config(self.config))
+
+        with patch.object(config_module, "save_config", return_value=False):
+            response = await config_routes.update_config_endpoint(_FakeRequest({
+                "workers": [{"id": "pool-1", "type": "local", "port": 8188}]
+            }))
+
+        self.assertEqual(response.status, 500)
+        self.assertEqual(self.secrets.get("pool-1"), "bulk-restore-key")
+        self.assertEqual(self.read_config_file()["workers"][0]["type"], "serverless")
+
+    async def test_secret_replacement_is_rolled_back_when_config_persistence_fails(self):
+        self.secrets.set("pool-1", "prior-key")
+        self.config["workers"] = [{
+            "id": "pool-1", "type": "serverless", "endpoint_id": "ep-123",
+        }]
+        self.assertTrue(config_module.save_config(self.config))
+
+        with patch.object(config_module, "save_config", return_value=False):
+            response = await config_routes.update_worker_endpoint(_FakeRequest({
+                "worker_id": "pool-1", "api_key": "replacement-key",
+            }))
+
+        self.assertEqual(response.status, 500)
+        self.assertEqual(self.secrets.get("pool-1"), "prior-key")
+        self.assertEqual(self.read_config_file()["workers"][0]["type"], "serverless")
+
+    async def test_delete_secret_is_rolled_back_when_config_persistence_fails(self):
+        self.secrets.set("pool-1", "delete-restore-key")
+        self.config["workers"] = [{
+            "id": "pool-1", "type": "serverless", "endpoint_id": "ep-123",
+        }]
+        self.assertTrue(config_module.save_config(self.config))
+
+        with patch.object(config_module, "save_config", return_value=False):
+            response = await config_routes.delete_worker_endpoint(_FakeRequest({"worker_id": "pool-1"}))
+
+        self.assertEqual(response.status, 500)
+        self.assertEqual(self.secrets.get("pool-1"), "delete-restore-key")
+        self.assertEqual(self.read_config_file()["workers"][0]["id"], "pool-1")
+
+    async def test_regular_worker_delete_does_not_require_posix_secret_storage(self):
+        self.config["workers"] = [{"id": "local-1", "type": "local", "port": 8188}]
+        self.assertTrue(config_module.save_config(self.config))
+        with patch.object(config_routes, "_serverless_secret_store", side_effect=RuntimeError("POSIX only")):
+            response = await config_routes.delete_worker_endpoint(_FakeRequest({"worker_id": "local-1"}))
+        self.assertEqual(response.status, 200)
+        self.assertEqual(self.read_config_file()["workers"], [])
+
+    async def test_enabling_serverless_without_key_is_rejected(self):
+        response = await config_routes.update_worker_endpoint(_FakeRequest({
+            "worker_id": "pool-1", "type": "serverless", "endpoint_id": "ep-123", "enabled": True,
+        }))
+        self.assertEqual(response.status, 400)
+        self.assertEqual(self.read_config_file()["workers"], [])
+
+    async def test_bulk_endpoint_change_is_blocked_during_active_run(self):
+        self.config["workers"] = [{"id": "pool-1", "type": "serverless", "endpoint_id": "ep-123"}]
+        self.assertTrue(config_module.save_config(self.config))
+        runtime_module.manager = types.SimpleNamespace(
+            tasks={"run-1": _RunningTask()},
+            runs={"run-1": {"worker": {"id": "serverless:pool-1:1", "pool_id": "pool-1"}}},
+            pool_is_active=lambda worker_id: worker_id == "pool-1",
+        )
+        updated = copy.deepcopy(self.config["workers"])
+        updated[0]["endpoint_id"] = "ep-other"
+        response = await config_routes.update_config_endpoint(_FakeRequest({"workers": updated}))
+        self.assertEqual(response.status, 409)
+        self.assertEqual(self.read_config_file()["workers"][0]["endpoint_id"], "ep-123")
+
+    async def test_non_serverless_update_behavior_is_preserved(self):
+        self.config["workers"] = [{"id": "local-1", "name": "Local", "type": "local", "port": 8188}]
+        self.assertTrue(config_module.save_config(self.config))
+
+        response = await config_routes.update_worker_endpoint(_FakeRequest({
+            "worker_id": "local-1", "name": "Renamed", "enabled": True,
+        }))
+
+        self.assertEqual(response.status, 200)
+        worker = self.read_config_file()["workers"][0]
+        self.assertEqual(worker["name"], "Renamed")
+        self.assertTrue(worker["enabled"])
+
+    async def test_config_normalization_removes_legacy_local_fields_from_serverless_records(self):
+        raw = {
+            "workers": [{
+                "id": "pool-1", "type": "serverless", "endpoint_id": "ep-123",
+                "host": "localhost", "port": 8188, "cuda_device": 2, "extra_args": "--old",
+            }]
+        }
+        normalized = config_module._merge_with_defaults(raw, config_module.get_default_config())
+        worker = normalized["workers"][0]
+        for obsolete in ("host", "port", "cuda_device", "extra_args"):
+            self.assertNotIn(obsolete, worker)
+
+
+class PersistedPoolGuardTests(_ConfigRouteFixture):
+    """Use the actual manager/ledger across restart, not a mocked idle result."""
+
+    def setUp(self):
+        super().setUp()
+        from test_serverless_runtime import runtime
+        from unittest.mock import AsyncMock
+        self.runtime = runtime
+        self.session_factory = AsyncMock()
+        self.ledger_path = self.root / "jobs.json"
+        self.manager = runtime.ServerlessManager(types.SimpleNamespace(), self.secrets, self.ledger_path, self.session_factory)
+        runtime_module.manager = self.manager
+        self.secrets.set("pool-1", "saved-key")
+        self.config["workers"] = [{"id": "pool-1", "type": "serverless", "endpoint_id": "ep-123", "concurrency": 1}]
+        self.assertTrue(config_module.save_config(self.config))
+
+    async def test_retained_jobs_block_destructive_changes_after_restart(self):
+        original = self.read_config_file()
+        for status in ("UNKNOWN", "CANCEL_REQUESTED", "IN_QUEUE", "IN_PROGRESS", "DOWNLOADING", "SUBMITTING"):
+            for restart in (False, True):
+                with self.subTest(status=status, restart=restart):
+                    self.manager.ledger.update("retained", worker_id="pool-1", endpoint_id="ep-123", status=status, job_id=None if status in {"UNKNOWN", "SUBMITTING"} else "owned-id")
+                    if restart:
+                        self.manager = self.runtime.ServerlessManager(types.SimpleNamespace(), self.secrets, self.ledger_path, self.session_factory)
+                        runtime_module.manager = self.manager
+                    else:
+                        self.manager.tasks["retained"] = types.SimpleNamespace(done=lambda: True)
+                    for fields in ({"api_key": "new-key"}, {"clear_api_key": True}, {"endpoint_id": "ep-other"}, {"type": "local", "port": 8188}):
+                        response = await config_routes.update_worker_endpoint(_FakeRequest(dict(fields, worker_id="pool-1")))
+                        self.assertEqual(response.status, 409, fields)
+                    response = await config_routes.delete_worker_endpoint(_FakeRequest({"worker_id": "pool-1"}))
+                    self.assertEqual(response.status, 409)
+                    worker = original["workers"][0]
+                    for workers in ([], [dict(worker, type="local")], [dict(worker, endpoint_id="ep-other")]):
+                        response = await config_routes.update_config_endpoint(_FakeRequest({"workers": workers}))
+                        self.assertEqual(response.status, 409, workers)
+                    self.assertEqual(self.read_config_file(), original)
+                    self.assertEqual(self.secrets.get("pool-1"), "saved-key")
+        self.session_factory.assert_not_awaited()
+
+    async def test_retained_jobs_allow_images_per_run_and_legacy_concurrency_saves(self):
+        # Upgrading from a config with concurrency=1 must not trip the busy guard,
+        # and the image count stays editable while old jobs are unresolved.
+        self.manager.ledger.update("retained", worker_id="pool-1", endpoint_id="ep-123", status="UNKNOWN")
+        response = await config_routes.update_worker_endpoint(_FakeRequest({
+            "worker_id": "pool-1", "endpoint_id": "ep-123", "jobs_per_run": 2, "concurrency": 2,
+        }))
+        self.assertEqual(response.status, 200)
+        saved = self.read_config_file()["workers"][0]
+        self.assertEqual(saved["jobs_per_run"], 2)
+        self.assertNotIn("concurrency", saved)
+        response = await config_routes.update_config_endpoint(_FakeRequest({"workers": [dict(saved, jobs_per_run=3, concurrency=5)]}))
+        self.assertEqual(response.status, 200)
+        self.assertEqual(self.read_config_file()["workers"][0]["jobs_per_run"], 3)
+
+    async def test_retained_job_allows_disable_and_rename_but_not_key_clear(self):
+        self.manager.ledger.update("retained", worker_id="pool-1", status="UNKNOWN")
+        response = await config_routes.update_worker_endpoint(_FakeRequest({"worker_id": "pool-1", "enabled": False, "name": "Paused"}))
+        self.assertEqual(response.status, 200)
+        self.assertEqual(self.read_config_file()["workers"][0]["name"], "Paused")
+        response = await config_routes.update_worker_endpoint(_FakeRequest({"worker_id": "pool-1", "clear_api_key": True}))
+        self.assertEqual(response.status, 409)
+        self.assertEqual(self.secrets.get("pool-1"), "saved-key")
+
+    async def test_terminal_or_other_pool_records_do_not_block_deletion(self):
+        self.manager.ledger.update("done", worker_id="pool-1", status="RESOLVED")
+        self.manager.ledger.update("other", worker_id="other", status="UNKNOWN")
+        response = await config_routes.delete_worker_endpoint(_FakeRequest({"worker_id": "pool-1"}))
+        self.assertEqual(response.status, 200)
+        self.assertIsNone(self.secrets.get("pool-1"))
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/test_serverless_runtime.py

@@ -0,0 +1,170 @@
+import asyncio
+import importlib.util
+import sys
+import tempfile
+import types
+import unittest
+from pathlib import Path
+from unittest.mock import AsyncMock, patch
+
+ROOT = Path(__file__).resolve().parents[1]
+PACKAGE = "dist_serverless_runtime_tests"
+for name, path in [(PACKAGE, ROOT), (PACKAGE + ".api", ROOT / "api"), (PACKAGE + ".api.orchestration", ROOT / "api/orchestration")]:
+    module = types.ModuleType(name)
+    module.__path__ = [str(path)]
+    sys.modules[name] = module
+
+def load(name):
+    spec = importlib.util.spec_from_file_location(PACKAGE + ".api.orchestration." + name, ROOT / "api/orchestration" / (name + ".py"))
+    assert spec and spec.loader
+    module = importlib.util.module_from_spec(spec)
+    sys.modules[spec.name] = module
+    spec.loader.exec_module(module)
+    return module
+
+runtime = load("serverless_runtime")
+
+class Secrets:
+    def get(self, worker_id):
+        return "private-test-key"
+
+class RuntimeTests(unittest.IsolatedAsyncioTestCase):
+    async def asyncSetUp(self):
+        self.temp = tempfile.TemporaryDirectory()
+        self.server = types.SimpleNamespace(distributed_pending_jobs={"collector-job": asyncio.Queue()}, distributed_jobs_lock=asyncio.Lock())
+        self.manager = runtime.ServerlessManager(self.server, Secrets(), Path(self.temp.name) / "jobs.json", AsyncMock(return_value=object()))
+        self.manager.poll_seconds = 0
+        self.worker = {"id": "pool:1", "pool_id": "pool", "endpoint_id": "ep1", "timeout_seconds": 30, "type": "serverless"}
+        self.prepared = {"worker": self.worker, "workflow": {"1": {}}, "images": [], "manifest": [{"job_id": "collector-job", "prefix": "dist_x", "kind": "image"}]}
+
+    async def asyncTearDown(self):
+        await self.manager.cancel()
+        self.temp.cleanup()
+
+    async def test_job_id_persisted_before_status_and_results_injected(self):
+        api = types.SimpleNamespace(submit=AsyncMock(return_value="provider-job"), cancel=AsyncMock())
+        async def status(job_id):
+            self.assertEqual(self.manager.ledger.all()[0]["job_id"], "provider-job")
+            return {"status": "COMPLETED", "output": {"images": []}}
+        api.status = status
+        with patch.object(runtime, "RunpodClient", return_value=api), patch.object(runtime, "decode_results", AsyncMock(return_value={"collector-job": ["tensor"]})):
+            await self.manager.start("execution", "prompt", [self.prepared])
+            await self.manager.wait_for_execution("execution")
+        self.assertEqual(api.submit.await_count, 1)
+        item = self.server.distributed_pending_jobs["collector-job"].get_nowait()
+        self.assertEqual(item["worker_id"], "pool:1")
+        self.assertEqual(item["tensor"], "tensor")
+        self.assertTrue(item["is_last"])
+        self.assertEqual(self.manager.ledger.all()[0]["status"], "COMPLETED")
+
+    async def test_output_error_marks_failed_and_delivers_explicit_error(self):
+        api = types.SimpleNamespace(submit=AsyncMock(return_value="provider-job"), status=AsyncMock(return_value={"status": "COMPLETED", "output": {"error": "private-test-key"}}), cancel=AsyncMock())
+        with patch.object(runtime, "RunpodClient", return_value=api):
+            await self.manager.start("execution", "prompt", [self.prepared])
+            await self.manager.wait_for_execution("execution")
+        record = self.manager.ledger.all()[0]
+        self.assertEqual(record["status"], "FAILED")
+        item = self.server.distributed_pending_jobs["collector-job"].get_nowait()
+        self.assertIn("serverless_error", item)
+        self.assertNotIn("private-test-key", str(record) + str(item))
+
+    async def test_ambiguous_submit_never_retries(self):
+        api = types.SimpleNamespace(submit=AsyncMock(side_effect=runtime.AmbiguousSubmission("Unknown submission")), cancel=AsyncMock())
+        with patch.object(runtime, "RunpodClient", return_value=api):
+            await self.manager.start("execution", "prompt", [self.prepared])
+            await self.manager.wait_for_execution("execution")
+        self.assertEqual(api.submit.await_count, 1)
+        self.assertEqual(self.manager.ledger.all()[0]["status"], "UNKNOWN")
+
+    async def test_cancel_targets_owned_job_only(self):
+        polled = asyncio.Event()
+        async def status(_job):
+            polled.set()
+            await asyncio.Event().wait()
+        api = types.SimpleNamespace(submit=AsyncMock(return_value="owned-job"), status=status, cancel=AsyncMock(return_value={"status": "CANCELLED"}))
+        with patch.object(runtime, "RunpodClient", return_value=api):
+            await self.manager.start("execution", "prompt", [self.prepared])
+            await polled.wait()
+            await self.manager.cancel(worker_id="other-pool")
+            api.cancel.assert_not_awaited()
+            await self.manager.cancel(worker_id="pool")
+        api.cancel.assert_awaited_once_with("owned-job")
+
+    async def test_all_images_of_a_run_submit_at_once_ignoring_legacy_concurrency(self):
+        self.assertEqual(runtime.PER_WORKER_JOB_LIMIT, 16)
+        gate = asyncio.Event()
+        submitted = []
+        async def submit(*args):
+            submitted.append(args)
+            return f"provider-{len(submitted)}"
+        async def status(_job):
+            await gate.wait()
+            return {"status": "COMPLETED", "output": {"images": []}}
+        api = types.SimpleNamespace(submit=submit, status=status, cancel=AsyncMock())
+        # A config saved before the setting was removed still carries concurrency=1.
+        legacy = dict(self.worker, concurrency=1)
+        jobs = [dict(self.prepared, worker=dict(legacy, id=f"pool:{i}")) for i in range(1, 4)]
+        with patch.object(runtime, "RunpodClient", return_value=api), patch.object(runtime, "decode_results", AsyncMock(return_value={"collector-job": ["tensor"]})):
+            await self.manager.start("execution", "prompt", jobs)
+            for _ in range(50):
+                if len(submitted) == 3:
+                    break
+                await asyncio.sleep(0)
+            self.assertEqual(len(submitted), 3)
+            gate.set()
+            await self.manager.wait_for_execution("execution")
+
+    async def test_per_pool_job_limit_bounds_submission(self):
+        patcher = patch.object(runtime, "PER_WORKER_JOB_LIMIT", 1)
+        patcher.start()
+        self.addCleanup(patcher.stop)
+        gate = asyncio.Event()
+        reached = asyncio.Event()
+        count = 0
+        async def submit(*args):
+            nonlocal count
+            count += 1
+            reached.set()
+            return f"provider-{count}"
+        async def status(_job):
+            await gate.wait()
+            return {"status": "COMPLETED", "output": {"images": []}}
+        api = types.SimpleNamespace(submit=submit, status=status, cancel=AsyncMock())
+        second = dict(self.prepared, worker=dict(self.worker, id="pool:2"))
+        with patch.object(runtime, "RunpodClient", return_value=api), patch.object(runtime, "decode_results", AsyncMock(return_value={"collector-job": ["tensor"]})):
+            await self.manager.start("execution", "prompt", [self.prepared, second])
+            await reached.wait()
+            await asyncio.sleep(0)
+            self.assertEqual(count, 1)
+            gate.set()
+            await self.manager.wait_for_execution("execution")
+        self.assertEqual(count, 2)
+
+    async def test_cancel_persisted_job_after_master_restart(self):
+        self.manager.ledger.update("old", worker_id="pool", endpoint_id="old-endpoint", job_id="old-job", status="IN_PROGRESS", created_at=1, updated_at=1)
+        api = types.SimpleNamespace(cancel=AsyncMock(return_value={"status": "CANCELLED"}))
+        with patch.object(runtime, "RunpodClient", return_value=api) as client:
+            await self.manager.cancel(worker_id="pool")
+        api.cancel.assert_awaited_once_with("old-job")
+        self.assertEqual(client.call_args.args[1], "old-endpoint")
+        self.assertEqual(self.manager.ledger.all()[0]["status"], "CANCELLED")
+
+    async def test_master_failure_cancels_job_without_partial_success(self):
+        self.server.prompt_queue = types.SimpleNamespace(get_history=lambda _id: {"prompt": {"status": {"status_str": "error"}}})
+        api = types.SimpleNamespace(submit=AsyncMock(return_value="job"), cancel=AsyncMock(return_value={"status": "CANCELLED"}))
+        with patch.object(runtime, "RunpodClient", return_value=api):
+            await self.manager.start("execution", "prompt", [self.prepared])
+            await self.manager.wait_for_execution("execution")
+        api.submit.assert_not_awaited()
+        self.assertEqual(self.manager.ledger.all()[0]["status"], "FAILED")
+
+    async def test_deadline_includes_local_slot_wait(self):
+        self.prepared["worker"] = dict(self.worker, timeout_seconds=0.01)
+        api = types.SimpleNamespace(submit=AsyncMock(return_value="job"), status=AsyncMock(side_effect=lambda _: asyncio.sleep(5)), cancel=AsyncMock(return_value={"status": "CANCELLED"}))
+        # A blocked slot expires without a billable submission.
+        self.manager.pool_limits["pool"] = asyncio.Semaphore(0)
+        with patch.object(runtime, "RunpodClient", return_value=api):
+            await self.manager.start("execution", "prompt", [self.prepared])
+            await self.manager.wait_for_execution("execution")
+        api.submit.assert_not_awaited()
+        self.assertEqual(self.manager.ledger.all()[0]["status"], "TIMED_OUT")

tests/test_serverless_status.py

@@ -0,0 +1,98 @@
+import importlib.util
+import json
+import sys
+import types
+import unittest
+from pathlib import Path
+from unittest.mock import AsyncMock, patch
+
+ROOT = Path(__file__).resolve().parents[1]
+PKG = "dist_serverless_status_tests"
+
+
+class StatusTests(unittest.IsolatedAsyncioTestCase):
+    async def asyncSetUp(self):
+        def mod(name, **kwargs):
+            result = types.ModuleType(name)
+            result.__dict__.update(kwargs)
+            return result
+        self.manager = types.SimpleNamespace(health=AsyncMock(return_value={"ok": True, "queueRemaining": 0, "message": "Ready (scales from zero)"}), public_jobs=lambda _: [], cancel=AsyncMock(return_value=1))
+        self.pool = {"id": "pool", "type": "serverless", "endpoint_id": "ep"}
+        modules = {PKG: mod(PKG, __path__=[str(ROOT)]), PKG + ".api": mod("api", __path__=[str(ROOT/'api')]),
+                   PKG + ".utils": mod("utils", __path__=[]),
+                   PKG + ".utils.config": mod("config", load_config=lambda: {"workers": [self.pool]}),
+                   PKG + ".api.orchestration": mod("orchestration", __path__=[str(ROOT/'api/orchestration')]),
+                   PKG + ".api.orchestration.serverless_runtime": mod("runtime", get_manager=lambda: self.manager),
+                   "server": mod("server", PromptServer=types.SimpleNamespace(instance=types.SimpleNamespace(routes=types.SimpleNamespace(get=lambda _: lambda f: f, post=lambda _: lambda f: f))))}
+        with patch.dict(sys.modules, modules):
+            spec = importlib.util.spec_from_file_location(PKG+'.api.serverless_routes', ROOT/'api/serverless_routes.py')
+            self.routes = importlib.util.module_from_spec(spec)
+            spec.loader.exec_module(self.routes)
+
+    async def test_health_only_never_submits_and_supports_scale_to_zero(self):
+        request = types.SimpleNamespace(json=AsyncMock(return_value={"worker_id": "pool"}))
+        result = await self.routes.test_serverless_endpoint(request)
+        self.assertEqual(result.status, 200)
+        self.assertTrue(json.loads(result.text)["ok"])
+        self.manager.health.assert_awaited_once_with(self.pool, fresh=True)
+        self.manager.cancel.assert_not_awaited()
+
+    async def test_unknown_pool_cannot_cancel_other_jobs(self):
+        request = types.SimpleNamespace(json=AsyncMock(return_value={"worker_id": "other"}))
+        result = await self.routes.cancel_serverless_endpoint(request)
+        self.assertEqual(result.status, 404)
+        self.manager.cancel.assert_not_awaited()
+
+    async def test_transport_error_does_not_leak_credentials(self):
+        self.manager.health.side_effect = RuntimeError("Bearer secret-token https://signed.invalid/?token=value")
+        request = types.SimpleNamespace(match_info={"worker_id": "pool"})
+        result = await self.routes.serverless_status_endpoint(request)
+        self.assertFalse(json.loads(result.text)["ok"])
+        self.assertNotIn("secret-token", result.text)
+        self.assertNotIn("signed.invalid", result.text)
+
+    async def test_cancel_uses_only_requested_pool(self):
+        request = types.SimpleNamespace(json=AsyncMock(return_value={"worker_id": "pool"}))
+        result = await self.routes.cancel_serverless_endpoint(request)
+        self.assertEqual(result.status, 200)
+        self.manager.cancel.assert_awaited_once_with(worker_id="pool")
+
+    async def test_reconcile_only_reads_the_requested_pools_jobs(self):
+        self.manager.reconcile = AsyncMock(return_value=2)
+        request = types.SimpleNamespace(json=AsyncMock(return_value={"worker_id": "pool"}))
+        result = await self.routes.reconcile_serverless_endpoint(request)
+        self.assertEqual(result.status, 200)
+        self.assertEqual(json.loads(result.text)["resolved"], 2)
+        self.manager.reconcile.assert_awaited_once_with(worker_id="pool")
+        self.manager.cancel.assert_not_awaited()
+
+    async def test_resolve_requires_explicit_boolean_confirmation_and_record_id(self):
+        self.manager.acknowledge_recovery = AsyncMock()
+        for payload in ({"worker_id": "pool", "record_id": "record"},
+                        {"worker_id": "pool", "record_id": "record", "provider_checked": "true"},
+                        {"worker_id": "pool", "provider_checked": True}):
+            with self.subTest(payload=payload):
+                result = await self.routes.resolve_serverless_endpoint(types.SimpleNamespace(json=AsyncMock(return_value=payload)))
+                self.assertEqual(result.status, 400)
+        self.manager.acknowledge_recovery.assert_not_awaited()
+
+    async def test_resolve_is_scoped_and_does_not_claim_provider_cancellation(self):
+        self.manager.acknowledge_recovery = AsyncMock()
+        request = types.SimpleNamespace(json=AsyncMock(return_value={"worker_id": "pool", "record_id": "record", "provider_checked": True}))
+        result = await self.routes.resolve_serverless_endpoint(request)
+        self.assertEqual(result.status, 200)
+        self.manager.acknowledge_recovery.assert_awaited_once_with("pool", "record", provider_checked=True)
+        self.assertIn("not provider-confirmed", json.loads(result.text)["message"])
+        self.manager.cancel.assert_not_awaited()
+
+    async def test_live_or_mismatched_record_refuses_manual_resolution(self):
+        self.manager.acknowledge_recovery = AsyncMock(side_effect=ValueError("Unresolved inactive job not found for this pool"))
+        request = types.SimpleNamespace(json=AsyncMock(return_value={"worker_id": "pool", "record_id": "record", "provider_checked": True}))
+        result = await self.routes.resolve_serverless_endpoint(request)
+        self.assertEqual(result.status, 409)
+
+    async def test_failed_health_still_exposes_recovery_ids(self):
+        self.manager.health.side_effect = RuntimeError("private transport error")
+        self.manager.public_jobs = lambda _: [{"record_id": "record", "status": "UNKNOWN"}]
+        result = await self.routes.serverless_status_endpoint(types.SimpleNamespace(match_info={"worker_id": "pool"}))
+        self.assertEqual(json.loads(result.text)["jobs"][0]["record_id"], "record")

tests/test_static_mode.py

@@ -4,13 +4,12 @@ import sys
 import types
 import unittest
 from pathlib import Path
-from types import SimpleNamespace
 
 import torch
 
 
 def _load_static_mode_module():
-    module_path = Path(__file__).resolve().parents[3] / "upscale" / "modes" / "static.py"
+    module_path = Path(__file__).resolve().parents[1] / "upscale" / "modes" / "static.py"
     package_name = "dist_static_mode_testpkg"
 
     for mod_name in list(sys.modules):
@@ -66,34 +65,8 @@ def _load_static_mode_module():
         arr = np.array(image).astype(np.float32) / 255.0
         return torch.from_numpy(arr).unsqueeze(0)
 
-    def _blend_processed_batch_item(
-        result_images,
-        processed_batch,
-        batch_index,
-        blend_fn,
-        x1,
-        y1,
-        ew,
-        eh,
-        tile_mask,
-        padding,
-    ):
-        tile_pil = _tensor_to_pil(processed_batch, batch_index)
-        if tile_pil.size != (ew, eh):
-            tile_pil = tile_pil.resize((ew, eh), PILImage.LANCZOS)
-        result_images[batch_index] = blend_fn(
-            result_images[batch_index],
-            tile_pil,
-            x1,
-            y1,
-            (ew, eh),
-            tile_mask,
-            padding,
-        )
-
     image_module.tensor_to_pil = _tensor_to_pil
     image_module.pil_to_tensor = _pil_to_tensor
-    image_module.blend_processed_batch_item = _blend_processed_batch_item
     sys.modules[f"{package_name}.utils.image"] = image_module
 
     async_helpers_module = types.ModuleType(f"{package_name}.utils.async_helpers")
@@ -129,10 +102,10 @@ def _load_static_mode_module():
         distributed_pending_tile_jobs={},
     )
     job_store_module.init_static_job_batched = _noop
-    job_store_module.mark_task_completed = _noop
-    job_store_module.cleanup_job = _noop
-    job_store_module.drain_results_queue = _noop
-    job_store_module.get_completed_count = _noop
+    job_store_module._mark_task_completed = _noop
+    job_store_module._cleanup_job = _noop
+    job_store_module._drain_results_queue = _noop
+    job_store_module._get_completed_count = _noop
     sys.modules[f"{package_name}.upscale.job_store"] = job_store_module
 
     job_models_module = types.ModuleType(f"{package_name}.upscale.job_models")
@@ -143,36 +116,6 @@ def _load_static_mode_module():
     job_models_module.TileJobState = _TileJobState
     sys.modules[f"{package_name}.upscale.job_models"] = job_models_module
 
-    processing_args_path = Path(__file__).resolve().parents[3] / "upscale" / "processing_args.py"
-    processing_args_spec = importlib.util.spec_from_file_location(
-        f"{package_name}.upscale.processing_args",
-        processing_args_path,
-    )
-    processing_args_module = importlib.util.module_from_spec(processing_args_spec)
-    assert processing_args_spec is not None and processing_args_spec.loader is not None
-    sys.modules[processing_args_spec.name] = processing_args_module
-    processing_args_spec.loader.exec_module(processing_args_module)
-
-    tile_processing_path = Path(__file__).resolve().parents[3] / "upscale" / "tile_processing.py"
-    tile_processing_spec = importlib.util.spec_from_file_location(
-        f"{package_name}.upscale.tile_processing",
-        tile_processing_path,
-    )
-    tile_processing_module = importlib.util.module_from_spec(tile_processing_spec)
-    assert tile_processing_spec is not None and tile_processing_spec.loader is not None
-    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
@@ -205,20 +148,14 @@ class _FakeStaticWorker(static_mode.StaticModeMixin):
     def _poll_job_ready(self, *_args, **_kwargs):
         return self.job_ready
 
-    async def request_assignment(self, *_args, **_kwargs):
+    async def _request_tile_from_master(self, *_args, **_kwargs):
         self.request_calls += 1
-        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,
-        )
+        return self.tile_sequence.pop(0)
 
-    async def send_heartbeat(self, *_args, **_kwargs):
+    async def _send_heartbeat_to_master(self, *_args, **_kwargs):
         self.heartbeat_calls += 1
 
-    async def send_tiles_batch(
+    async def send_tiles_batch_to_master(
         self,
         processed_tiles,
         _multi_job_id,

tests/test_transfer_codec.py

@@ -0,0 +1,219 @@
+import importlib.util
+import io
+import sys
+from pathlib import Path
+
+import numpy as np
+import pytest
+import torch
+from PIL import Image
+
+
+def _load_codec():
+    path = Path(__file__).resolve().parents[1] / "utils" / "transfer_codec.py"
+    spec = importlib.util.spec_from_file_location("dist_transfer_codec_test", path)
+    module = importlib.util.module_from_spec(spec)
+    sys.modules["dist_transfer_codec_test"] = module
+    spec.loader.exec_module(module)
+    return module
+
+
+tc = _load_codec()
+
+
+def _smooth_frames(count=4, height=67, width=91):
+    """Smooth moving gradient, closer to generated content than random noise."""
+    yy, xx = np.mgrid[0:height, 0:width].astype(np.float32)
+    frames = []
+    for i in range(count):
+        r = 0.5 + 0.5 * np.sin((xx + 3 * i) / 13.0)
+        g = 0.5 + 0.5 * np.cos((yy - 2 * i) / 11.0)
+        b = (xx + yy + 5 * i) / (width + height + 5 * count)
+        frames.append(np.stack([r, g, b], axis=-1))
+    return torch.from_numpy(np.stack(frames)).float()
+
+
+def _legacy_master_tensor(frame_tensor):
+    """What the unmodified pipeline delivered: tensor_to_pil -> PNG -> pil_to_tensor."""
+    img = Image.fromarray((255 * frame_tensor.numpy()).astype(np.uint8))
+    buf = io.BytesIO()
+    img.save(buf, format="PNG", compress_level=0)
+    with Image.open(io.BytesIO(buf.getvalue())) as decoded:
+        return np.array(decoded.convert("RGB")).astype(np.float32) / 255.0
+
+
+def _round_trip(frames_u8, fmt, quality=90, max_chunk_bytes=10**9):
+    out = []
+    chunks = list(tc.encode_chunks(frames_u8, fmt, quality, max_chunk_bytes))
+    for chunk in chunks:
+        envelope = {k: v for k, v in chunk.items() if k != "parts"}
+        out.extend(tc.decode_chunk(envelope, chunk["parts"]))
+    return chunks, out
+
+
+def _psnr(a, b):
+    mse = np.mean((a.astype(np.float64) - b.astype(np.float64)) ** 2)
+    return float("inf") if mse == 0 else 10 * np.log10(255.0**2 / mse)
+
+
+@pytest.mark.parametrize("fmt", ["png", "ffv1"])
+def test_lossless_formats_reproduce_legacy_master_tensors_exactly(fmt):
+    images = _smooth_frames()
+    frames = tc.tensor_batch_to_uint8(images)
+    _, decoded = _round_trip(frames, fmt)
+
+    assert len(decoded) == images.shape[0]
+    for i, frame in enumerate(decoded):
+        got = tc.uint8_frame_to_tensor_array(frame)[0]
+        assert np.array_equal(got, _legacy_master_tensor(images[i]))
+
+
+@pytest.mark.parametrize("fmt", ["jpeg", "webp", "h264", "h265", "av1"])
+def test_lossy_formats_round_trip_odd_sizes_with_reasonable_quality(fmt):
+    if fmt in tc.VIDEO_FORMATS and not tc.video_codec_available(fmt):
+        pytest.skip(f"{fmt} encoder not available in this PyAV build")
+    frames = tc.tensor_batch_to_uint8(_smooth_frames())
+    _, decoded = _round_trip(frames, fmt, quality=90)
+
+    assert len(decoded) == frames.shape[0]
+    for i, frame in enumerate(decoded):
+        assert frame.shape == frames[i].shape
+        assert frame.dtype == np.uint8
+        assert _psnr(frame, frames[i]) > 30
+
+
+def test_h265_handles_frames_smaller_than_ctu():
+    if not tc.video_codec_available("h265"):
+        pytest.skip("h265 encoder not available")
+    frames = tc.tensor_batch_to_uint8(_smooth_frames(count=3, height=17, width=23))
+    _, decoded = _round_trip(frames, "h265")
+    assert [f.shape for f in decoded] == [(17, 23, 3)] * 3
+
+
+def test_image_formats_split_chunks_by_byte_cap_and_keep_order():
+    frames = tc.tensor_batch_to_uint8(_smooth_frames(count=5))
+    one = len(tc.encode_image_frame(frames[0], "png", 90))
+    chunks, decoded = _round_trip(frames, "png", max_chunk_bytes=one * 2 + 1)
+
+    assert len(chunks) >= 3
+    assert [c["start_index"] for c in chunks] == sorted(c["start_index"] for c in chunks)
+    assert sum(c["count"] for c in chunks) == 5
+    assert all(len(c["parts"]) == c["count"] for c in chunks)
+    assert np.array_equal(np.stack(decoded), frames)
+
+
+def test_video_formats_split_chunks_by_frame_cap(monkeypatch):
+    monkeypatch.setattr(tc, "MAX_VIDEO_CHUNK_FRAMES", 2)
+    frames = tc.tensor_batch_to_uint8(_smooth_frames(count=5))
+    chunks, decoded = _round_trip(frames, "ffv1")
+
+    assert [(c["start_index"], c["count"]) for c in chunks] == [(0, 2), (2, 2), (4, 1)]
+    assert all(len(c["parts"]) == 1 for c in chunks)
+    assert np.array_equal(np.stack(decoded), frames)
+
+
+def test_tensor_batch_to_uint8_clips_and_normalizes_channels():
+    rgba = torch.tensor([[[[1.5, -0.2, 0.5, 0.3]]]])
+    assert tc.tensor_batch_to_uint8(rgba).tolist() == [[[[255, 0, 127]]]]
+
+    gray = torch.full((1, 2, 2, 1), 0.5)
+    out = tc.tensor_batch_to_uint8(gray)
+    assert out.shape == (1, 2, 2, 3)
+    assert (out == 127).all()
+
+
+def test_quality_to_crf_mapping():
+    assert tc.quality_to_crf(100, 51) == 0
+    assert tc.quality_to_crf(90, 51) == 13
+    assert tc.quality_to_crf(80, 51) == 19
+    assert tc.quality_to_crf(90, 63) == 16
+    assert tc.quality_to_crf(1, 51) == 51
+    assert tc.quality_to_crf("bogus", 51) == tc.quality_to_crf(tc.DEFAULT_TRANSFER_QUALITY, 51)
+
+
+def test_normalize_format_rejects_unknown_values():
+    assert tc.normalize_format(" H264 ") == "h264"
+    assert tc.normalize_format(None) == tc.DEFAULT_TRANSFER_FORMAT
+    with pytest.raises(tc.TransferCodecError):
+        tc.normalize_format("gif")
+
+
+def test_decode_chunk_validates_envelope_against_payload():
+    frames = tc.tensor_batch_to_uint8(_smooth_frames(count=2))
+    chunk = next(tc.encode_chunks(frames, "png", 90, 10**9))
+    envelope = {k: v for k, v in chunk.items() if k != "parts"}
+
+    with pytest.raises(tc.TransferCodecError, match="image parts"):
+        tc.decode_chunk({**envelope, "count": 3}, chunk["parts"])
+    with pytest.raises(tc.TransferCodecError, match="expected"):
+        tc.decode_chunk({**envelope, "width": 10}, chunk["parts"])
+    with pytest.raises(tc.TransferCodecError, match="count"):
+        tc.decode_chunk({**envelope, "count": 0}, chunk["parts"])
+    with pytest.raises(tc.TransferCodecError):
+        tc.decode_chunk(envelope, [b"not an image", b"nope"])
+
+    video = next(tc.encode_chunks(frames, "ffv1", 90, 10**9))
+    video_env = {k: v for k, v in video.items() if k != "parts"}
+    with pytest.raises(tc.TransferCodecError, match="frames"):
+        tc.decode_chunk({**video_env, "count": 5}, video["parts"])
+
+
+def test_legacy_format_is_not_chunk_encoded():
+    frames = tc.tensor_batch_to_uint8(_smooth_frames(count=1))
+    with pytest.raises(tc.TransferCodecError):
+        list(tc.encode_chunks(frames, "legacy_png", 90, 10**9))
+
+
+# --------------------------------------------------------------------------- #
+# Size bounds on finalized payloads
+# --------------------------------------------------------------------------- #
+
+def _noise_frames(count=30, height=360, width=640, seed=97):
+    """High-entropy frames: worst case for every codec, and enough encoder
+    lookahead that pre-flush buffer sizes badly underestimate final sizes."""
+    return np.random.default_rng(seed).integers(0, 256, size=(count, height, width, 3), dtype=np.uint8)
+
+
+@pytest.mark.parametrize("fmt", ["h264", "ffv1"])
+def test_video_chunks_are_bounded_after_encoder_flush(fmt):
+    frames = _noise_frames()
+    limit = 1024 * 1024
+    # Sanity: one unconstrained encode of this batch is far above the limit,
+    # so the bound below is actually being enforced.
+    unbounded = next(tc.encode_chunks(frames, fmt, 90, 10**12))
+    assert len(unbounded["parts"][0]) > 3 * limit
+
+    chunks = list(tc.encode_chunks(frames, fmt, 90, limit))
+
+    assert all(len(c["parts"]) == 1 and len(c["parts"][0]) <= limit for c in chunks)
+    assert [c["start_index"] for c in chunks] == list(np.cumsum([0] + [c["count"] for c in chunks[:-1]]))
+    assert sum(c["count"] for c in chunks) == len(frames)
+    decoded = [f for c in chunks for f in tc.decode_chunk(c, c["parts"])]
+    assert len(decoded) == len(frames)
+    if fmt == "ffv1":
+        assert np.array_equal(np.stack(decoded), frames)
+
+
+def test_single_video_frame_may_exceed_chunk_target_up_to_frame_limit():
+    frames = _noise_frames(count=3)
+    one_frame = len(next(tc.encode_chunks(frames[:1], "ffv1", 90, 10**12))["parts"][0])
+
+    chunks = list(tc.encode_chunks(frames, "ffv1", 90, one_frame // 2, one_frame * 2))
+
+    assert [c["count"] for c in chunks] == [1, 1, 1]
+    assert all(one_frame // 2 < len(c["parts"][0]) <= one_frame * 2 for c in chunks)
+
+
+@pytest.mark.parametrize("fmt", ["ffv1", "h264", "png"])
+def test_frame_that_cannot_fit_alone_raises(fmt):
+    frames = _noise_frames(count=2)
+    with pytest.raises(tc.FrameTooLargeError, match="request limit"):
+        list(tc.encode_chunks(frames, fmt, 90, 64 * 1024, 64 * 1024))
+
+
+def test_still_chunks_are_bounded():
+    frames = _noise_frames(count=10)
+    limit = 1024 * 1024
+    chunks = list(tc.encode_chunks(frames, "jpeg", 90, limit))
+    assert all(sum(len(p) for p in c["parts"]) <= limit for c in chunks)
+    assert sum(c["count"] for c in chunks) == 10

tests/test_v3_collector_transfer.py

@@ -0,0 +1,125 @@
+"""Registered V3 DistributedCollector: transfer options in the public schema and forwarded to the runtime.
+
+Runs without a ComfyUI checkout by loading nodes/v3.py against a recording
+stand-in for comfy_api's io module. tests/test_v3_migration.py covers the same
+boundary against the real framework when COMFYUI_SOURCE_ROOT is set.
+"""
+import importlib.util
+import sys
+import types
+from pathlib import Path
+
+import pytest
+
+ROOT = Path(__file__).resolve().parents[1]
+PKG = "dist_v3_transfer_testpkg"
+
+
+class _Field:
+    def __init__(self, kind, id=None, **kwargs):
+        self.kind = kind
+        self.id = id
+        self.kwargs = kwargs
+
+
+def _io_type(kind):
+    return types.SimpleNamespace(
+        Input=lambda id, **kw: _Field(kind, id, **kw),
+        Output=lambda id=None, **kw: _Field(kind, id, **kw),
+    )
+
+
+def _fake_io():
+    io = types.ModuleType("io")
+
+    class ComfyNode:
+        pass
+
+    io.ComfyNode = ComfyNode
+    io.Schema = lambda **kw: types.SimpleNamespace(**kw)
+    io.NodeOutput = lambda *values, **kw: ("NodeOutput", values, kw)
+    io.NumberDisplay = types.SimpleNamespace(number="number")
+    io.Hidden = types.SimpleNamespace(unique_id="unique_id", extra_pnginfo="extra_pnginfo")
+    for name in ("Int", "Float", "String", "Boolean", "Combo", "Image", "Audio",
+                 "AnyType", "Model", "Conditioning", "Vae"):
+        setattr(io, name, _io_type(name))
+    return io
+
+
+@pytest.fixture
+def v3(monkeypatch):
+    calls = []
+
+    class CollectorRuntime:
+        def run(self, **kwargs):
+            calls.append(kwargs)
+            return (kwargs["images"], kwargs["audio"])
+
+    def module(name, **attrs):
+        mod = types.ModuleType(name)
+        mod.__dict__.update(attrs)
+        monkeypatch.setitem(sys.modules, name, mod)
+        return mod
+
+    for name in [n for n in sys.modules if n == PKG or n.startswith(PKG + ".")]:
+        monkeypatch.delitem(sys.modules, name)
+    module(PKG, __path__=[])
+    module(f"{PKG}.nodes", __path__=[])
+    module(f"{PKG}.utils", __path__=[])
+    module(f"{PKG}.nodes.utilities")
+    module(f"{PKG}.nodes.distributed_upscale", UltimateSDUpscaleDistributed=object)
+    module(f"{PKG}.nodes.collector", DistributedCollectorNode=CollectorRuntime)
+    module("comfy", __path__=[])
+    module("comfy.samplers", KSampler=types.SimpleNamespace(SAMPLERS=["euler"], SCHEDULERS=["normal"]))
+    module("comfy_api", __path__=[])
+    module("comfy_api.v0_0_2", io=_fake_io())
+
+    codec_spec = importlib.util.spec_from_file_location(f"{PKG}.utils.transfer_codec", ROOT / "utils" / "transfer_codec.py")
+    codec = importlib.util.module_from_spec(codec_spec)
+    monkeypatch.setitem(sys.modules, f"{PKG}.utils.transfer_codec", codec)
+    codec_spec.loader.exec_module(codec)
+
+    spec = importlib.util.spec_from_file_location(f"{PKG}.nodes.v3", ROOT / "nodes" / "v3.py")
+    mod = importlib.util.module_from_spec(spec)
+    monkeypatch.setitem(sys.modules, f"{PKG}.nodes.v3", mod)
+    spec.loader.exec_module(mod)
+    return mod, codec, calls
+
+
+def test_registered_collector_schema_exposes_optional_transfer_inputs(v3):
+    mod, codec, _ = v3
+    schema = mod.DistributedCollector.define_schema()
+    inputs = {field.id: field for field in schema.inputs}
+
+    assert [field.id for field in schema.inputs] == [
+        "load_balance", "images", "audio", "transfer_format", "transfer_quality"]
+    fmt = inputs["transfer_format"]
+    assert fmt.kind == "Combo"
+    assert fmt.kwargs["options"] == list(codec.TRANSFER_FORMATS)
+    assert fmt.kwargs["default"] == "legacy_png"
+    assert fmt.kwargs["optional"] is True
+    quality = inputs["transfer_quality"]
+    assert quality.kind == "Int"
+    assert (quality.kwargs["default"], quality.kwargs["min"], quality.kwargs["max"]) == (90, 1, 100)
+    assert quality.kwargs["optional"] is True
+    assert schema.is_input_list is True
+
+
+def test_registered_collector_forwards_transfer_options(v3):
+    mod, _, calls = v3
+
+    mod.DistributedCollector.execute(images=["img"], audio=[None], transfer_format=["av1"], transfer_quality=[70],
+                                     multi_job_id=["job"], is_worker=[True])
+
+    assert calls[-1]["transfer_format"] == ["av1"]
+    assert calls[-1]["transfer_quality"] == [70]
+    assert calls[-1]["multi_job_id"] == ["job"]
+
+
+def test_registered_collector_defaults_to_legacy_transfer(v3):
+    mod, _, calls = v3
+
+    mod.DistributedCollector.execute(images=["img"])
+
+    assert calls[-1]["transfer_format"] == "legacy_png"
+    assert calls[-1]["transfer_quality"] == 90

tests/test_v3_migration.py

@@ -0,0 +1,21 @@
+"""Real-framework acceptance; opt in with COMFYUI_SOURCE_ROOT."""
+import os
+from pathlib import Path
+import subprocess
+import sys
+
+import pytest
+
+
+def test_native_v3_runtime():
+    comfy_root = os.environ.get('COMFYUI_SOURCE_ROOT')
+    if not comfy_root:
+        pytest.skip('Set COMFYUI_SOURCE_ROOT to run native ComfyUI V3 acceptance')
+    helper = Path(__file__).parent / 'helpers' / 'v3_runtime_check.py'
+    result = subprocess.run(
+        [sys.executable, str(helper), str(Path(comfy_root).resolve())],
+        capture_output=True, text=True, timeout=120,
+        env={**os.environ, 'COMFYUI_IS_WORKER': '1'},
+    )
+    assert result.returncode == 0, result.stdout + '\n' + result.stderr
+    assert 'V3_ACCEPTANCE_OK' in result.stdout

tests/test_worker_ports.py

@@ -0,0 +1,113 @@
+import importlib.util
+import os
+import socket
+import unittest
+from pathlib import Path
+from unittest.mock import patch
+
+
+def load_ports():
+    path = Path(__file__).resolve().parents[1] / "workers" / "ports.py"
+    spec = importlib.util.spec_from_file_location("worker_ports_test", path)
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    return module
+
+
+class WorkerPortTests(unittest.TestCase):
+    def setUp(self):
+        self.ports = load_ports()
+
+    def test_allocation_starts_above_actual_master_and_skips_assigned_and_occupied(self):
+        workers = [
+            {"id": "existing", "host": "localhost", "port": 8190, "enabled": False},
+            {"id": "remote", "host": "remote.example", "port": 8192},
+        ]
+        with patch.object(self.ports, "is_port_available", side_effect=lambda port: port != 8191):
+            self.assertEqual(self.ports.allocate_worker_ports(8189, workers, 3), [8192, 8193, 8194])
+
+    def test_local_host_forms_reserve_ports(self):
+        for host in ["::1", "[::1]", "http://localhost/", "HTTPS://LOCALHOST", "0.0.0.0", None]:
+            with self.subTest(host=host), patch.object(self.ports, "is_port_available", return_value=True):
+                self.assertEqual(self.ports.allocate_worker_ports(8189, [{"id": "local", "host": host, "port": 8190}], 1), [8191])
+
+    def test_exhaustion_is_explicit_not_a_partial_allocation(self):
+        with patch.object(self.ports, "is_port_available", return_value=True):
+            with self.assertRaisesRegex(ValueError, "available.*ports"):
+                self.ports.allocate_worker_ports(65534, [], 2)
+
+    def test_launch_rejects_master_port_without_reassigning_configuration(self):
+        worker = {"id": "manual", "port": 8189}
+        with self.assertRaisesRegex(ValueError, "master.*8189"):
+            self.ports.validate_worker_port(worker, 8189, [])
+        self.assertEqual(worker["port"], 8189)
+
+    def test_launch_rejects_another_local_workers_reserved_port(self):
+        worker = {"id": "manual", "port": 8190}
+        other = {"id": "other", "host": None, "port": 8190, "enabled": False}
+        with self.assertRaisesRegex(ValueError, "other"):
+            self.ports.validate_worker_port(worker, 8189, [worker, other])
+
+    def test_launch_ignores_itself_and_remote_workers_with_same_port(self):
+        worker = {"id": "manual", "port": 8190}
+        other = {"id": "remote", "host": "example.com", "port": 8190}
+        with patch.object(self.ports, "is_port_available", return_value=True):
+            self.ports.validate_worker_port(worker, 8189, [worker, other])
+
+    def test_launch_rejects_an_occupied_port(self):
+        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as listener:
+            listener.bind(("127.0.0.1", 0))
+            listener.listen()
+            port = listener.getsockname()[1]
+            self.assertFalse(self.ports.is_port_available(port))
+            with self.assertRaisesRegex(ValueError, "already in use"):
+                self.ports.validate_worker_port({"id": "manual", "port": port}, 8189, [])
+
+    @unittest.skipIf(os.name == "nt", "asyncio does not reuse addresses on Windows")
+    def test_recently_closed_connection_does_not_block_worker_restart(self):
+        with socket.socket() as listener:
+            listener.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+            listener.bind(("127.0.0.1", 0))
+            listener.listen(1)
+            port = listener.getsockname()[1]
+            with socket.create_connection(("127.0.0.1", port), timeout=2) as client:
+                accepted, _ = listener.accept()
+                with accepted:
+                    accepted.settimeout(2)
+                    accepted.shutdown(socket.SHUT_WR)
+                    self.assertEqual(client.recv(1), b"")
+                    client.shutdown(socket.SHUT_WR)
+                    self.assertEqual(accepted.recv(1), b"")
+        self.assertTrue(self.ports.is_port_available(port))
+        self.ports.validate_worker_port({"id": "restart", "port": port}, 1, [])
+        with socket.socket() as restarted:
+            restarted.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+            restarted.bind(("127.0.0.1", port))
+            restarted.listen(1)
+
+    def test_active_listener_with_reuseaddr_is_still_a_conflict(self):
+        with socket.socket() as listener:
+            listener.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+            listener.bind(("127.0.0.1", 0))
+            listener.listen(1)
+            self.assertFalse(self.ports.is_port_available(listener.getsockname()[1]))
+
+    @unittest.skipUnless(socket.has_ipv6, "IPv6 unavailable")
+    def test_ipv6_only_listener_is_still_a_conflict(self):
+        with socket.socket(socket.AF_INET6, socket.SOCK_STREAM) as listener:
+            listener.setsockopt(socket.IPPROTO_IPV6, socket.IPV6_V6ONLY, 1)
+            try:
+                listener.bind(("::1", 0))
+            except OSError as exc:
+                self.skipTest(f"IPv6 loopback unavailable: {exc}")
+            listener.listen(1)
+            self.assertFalse(self.ports.is_port_available(listener.getsockname()[1]))
+
+    def test_invalid_ports_fail_clearly(self):
+        for port in [0, 65536, "not-a-port", None, True, 8190.5]:
+            with self.subTest(port=port), self.assertRaisesRegex(ValueError, "port"):
+                self.ports.validate_worker_port({"id": "manual", "port": port}, 8189, [])
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/test_worker_process_runtime.py

@@ -0,0 +1,352 @@
+import importlib.util
+import os
+import sys
+import tempfile
+import types
+import unittest
+from argparse import Namespace
+from pathlib import Path
+from unittest.mock import patch
+
+from sqlalchemy import create_engine
+from sqlalchemy.engine import URL, make_url
+
+
+def _load_process_module(module_filename: str):
+    module_path = Path(__file__).resolve().parents[1] / "workers" / "process" / module_filename
+    package_name = "dist_proc_testpkg"
+    module_name = module_filename[:-3]
+
+    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
+
+    workers_pkg = types.ModuleType(f"{package_name}.workers")
+    workers_pkg.__path__ = [str(module_path.parents[1])]
+    sys.modules[f"{package_name}.workers"] = workers_pkg
+
+    process_pkg = types.ModuleType(f"{package_name}.workers.process")
+    process_pkg.__path__ = []
+    sys.modules[f"{package_name}.workers.process"] = process_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
+
+    process_module = types.ModuleType(f"{package_name}.utils.process")
+    process_module.get_python_executable = lambda: "/usr/bin/test-python"
+    process_module.is_process_alive = lambda _pid: False
+    process_module.terminate_process = lambda *_args, **_kwargs: None
+    sys.modules[f"{package_name}.utils.process"] = process_module
+
+    config_module = types.ModuleType(f"{package_name}.utils.config")
+    config_module.load_config = lambda: {"workers": [], "settings": {"stop_workers_on_master_exit": False}}
+    config_module.save_config = lambda _config: None
+    sys.modules[f"{package_name}.utils.config"] = config_module
+    constants_module = types.ModuleType(f"{package_name}.utils.constants")
+    constants_module.PROCESS_TERMINATION_TIMEOUT = 1
+    constants_module.PROCESS_WAIT_TIMEOUT = 1
+    constants_module.WORKER_CHECK_INTERVAL = 0.01
+    sys.modules[f"{package_name}.utils.constants"] = constants_module
+    network_module = types.ModuleType(f"{package_name}.utils.network")
+    network_module.get_server_port = lambda: 8189
+    sys.modules[f"{package_name}.utils.network"] = network_module
+
+    spec = importlib.util.spec_from_file_location(
+        f"{package_name}.workers.process.{module_name}",
+        module_path,
+    )
+    module = importlib.util.module_from_spec(spec)
+    assert spec is not None and spec.loader is not None
+    spec.loader.exec_module(module)
+    return module
+
+
+root_discovery_module = _load_process_module("root_discovery.py")
+launch_builder_module = _load_process_module("launch_builder.py")
+
+
+class ComfyRootDiscoveryTests(unittest.TestCase):
+    def test_prefers_loaded_comfyui_module_path(self):
+        discovery = root_discovery_module.ComfyRootDiscovery()
+        server_module = types.SimpleNamespace(__file__="/opt/ComfyUI/server.py")
+
+        def fake_exists(path):
+            return path == "/opt/ComfyUI/main.py"
+
+        with patch.dict(sys.modules, {"server": server_module}, clear=False), \
+             patch.object(root_discovery_module.os.path, "exists", side_effect=fake_exists), \
+             patch.dict(root_discovery_module.os.environ, {}, clear=True):
+            self.assertEqual(discovery.find_comfy_root(), "/opt/ComfyUI")
+
+
+class LaunchCommandBuilderTests(unittest.TestCase):
+    def build_command(self, root, worker=None, runtime=None):
+        builder = launch_builder_module.LaunchCommandBuilder()
+        worker = worker or {"id": "worker-a", "port": 8190}
+        runtime = runtime if runtime is not None else Namespace(database_url=None)
+        with patch.object(builder, "_get_runtime_args", return_value=runtime):
+            return builder.build_launch_command(worker, str(root))
+
+    def test_database_is_unique_and_stable_by_worker_id(self):
+        with tempfile.TemporaryDirectory() as directory:
+            root = Path(directory)
+            (root / "main.py").touch()
+            first = self.build_command(root, {"id": "worker-a", "port": 8190})
+            renamed = self.build_command(root, {
+                "id": "worker-a", "port": 9001, "name": "Renamed", "cuda_device": 3,
+            })
+            second = self.build_command(root, {"id": "worker-b", "port": 8191})
+            database = first[first.index("--database-url") + 1]
+            self.assertEqual(database, renamed[renamed.index("--database-url") + 1])
+            self.assertNotEqual(database, second[second.index("--database-url") + 1])
+            self.assertTrue(database.startswith("sqlite:///" + root.as_posix() + "/"))
+            self.assertTrue(Path(database.removeprefix("sqlite:///")).parent.is_dir())
+            self.assertNotIn("--disable-assets", first)
+
+    def test_database_uses_effective_user_or_base_directory(self):
+        with tempfile.TemporaryDirectory() as directory:
+            root = Path(directory)
+            (root / "main.py").touch()
+            for extra_args, expected in [
+                (f'--user-directory "{root / "custom user"}"', root / "custom user"),
+                (f'--base-directory "{root / "custom base"}"', root / "custom base" / "user"),
+            ]:
+                with self.subTest(extra_args=extra_args):
+                    cmd = self.build_command(root, {
+                        "id": "worker-a", "port": 8190, "extra_args": extra_args,
+                    })
+                    database = cmd[cmd.index("--database-url") + 1]
+                    self.assertTrue(database.startswith("sqlite:///" + expected.as_posix() + "/"))
+            cmd = self.build_command(root, runtime=Namespace(
+                database_url="sqlite:///master.db", user_directory=str(root / "runtime user"),
+            ))
+            database = cmd[cmd.index("--database-url") + 1]
+            self.assertTrue(database.startswith("sqlite:///" + (root / "runtime user").as_posix() + "/"))
+            self.assertNotEqual(database, "sqlite:///master.db")
+
+    def test_relative_database_directories_use_worker_cwd(self):
+        with tempfile.TemporaryDirectory() as directory:
+            root = Path(directory) / "ComfyUI"
+            root.mkdir()
+            (root / "main.py").touch()
+            launcher = Path(directory) / "launcher"
+            launcher.mkdir()
+            cases = [
+                (Namespace(database_url=None, user_directory="profiles"), "", root / "profiles"),
+                (Namespace(database_url=None, base_directory="data"), "", root / "data/user"),
+                (Namespace(database_url=None), "--user-directory=profiles", root / "profiles"),
+                (Namespace(database_url=None), "--base-directory data", root / "data/user"),
+            ]
+            for runtime, extra, expected in cases:
+                with self.subTest(runtime=runtime, extra=extra), \
+                     patch.object(os, "getcwd", return_value=str(launcher)):
+                    cmd = self.build_command(root, {
+                        "id": "worker-a", "port": 8190, "extra_args": extra,
+                    }, runtime)
+                    database = make_url(cmd[cmd.index("--database-url") + 1]).database
+                    self.assertEqual(Path(database).parent, expected / "distributed/workers")
+            self.assertFalse((launcher / "profiles").exists())
+            self.assertFalse((launcher / "data").exists())
+
+    def test_database_url_round_trips_special_characters_without_collisions(self):
+        with tempfile.TemporaryDirectory() as directory:
+            root = Path(directory)
+            (root / "main.py").touch()
+            for name in ["user?profile", "user%3Fprofile"]:
+                with self.subTest(name=name):
+                    user = root / name
+                    runtime = Namespace(database_url=None, user_directory=str(user))
+                    # SQLAlchemy 2.0 cannot round-trip '?' in a filename. Fail
+                    # clearly on those versions rather than silently sharing a DB.
+                    sample = (user / "test.db").as_posix()
+                    serialized = URL.create("sqlite", database=sample).render_as_string()
+                    if make_url(serialized).database != sample:
+                        with self.assertRaisesRegex(ValueError, "SQLAlchemy.*database path"):
+                            self.build_command(root, runtime=runtime)
+                        self.assertFalse(user.exists())
+                        continue
+                    databases = []
+                    for worker_id in ["worker-a", "worker-b"]:
+                        cmd = self.build_command(root, {"id": worker_id, "port": 8190}, runtime)
+                        url = cmd[cmd.index("--database-url") + 1]
+                        database = make_url(url).database
+                        self.assertEqual(Path(database).parent, user / "distributed/workers")
+                        engine = create_engine(url)
+                        try:
+                            with engine.connect() as connection:
+                                actual = connection.exec_driver_sql("PRAGMA database_list").one()[2]
+                                self.assertEqual(Path(actual), Path(database))
+                        finally:
+                            engine.dispose()
+                        databases.append(database)
+                    self.assertNotEqual(*databases)
+
+    def test_explicit_database_and_disabled_assets_are_preserved(self):
+        with tempfile.TemporaryDirectory() as directory:
+            root = Path(directory)
+            (root / "main.py").touch()
+            for extra in ["--database-url sqlite:///explicit.db", "--database-url=sqlite:///explicit.db", "--disable-assets"]:
+                with self.subTest(extra=extra):
+                    cmd = self.build_command(root, {"id": "worker-a", "port": 8190, "extra_args": extra})
+                    self.assertEqual(sum(arg.split("=")[0] == "--database-url" for arg in cmd), 0 if extra == "--disable-assets" else 1)
+                    self.assertFalse((root / "user").exists())
+
+    def test_older_comfyui_without_database_flag_keeps_existing_command(self):
+        with tempfile.TemporaryDirectory() as directory:
+            root = Path(directory)
+            (root / "main.py").touch()
+            cmd = self.build_command(root, runtime=Namespace())
+            self.assertNotIn("--database-url", cmd)
+            self.assertFalse((root / "user").exists())
+
+    def test_worker_id_cannot_escape_database_directory(self):
+        with tempfile.TemporaryDirectory() as directory:
+            root = Path(directory)
+            (root / "main.py").touch()
+            databases = []
+            for worker_id in ["../outside", "a/b", "a_b"]:
+                cmd = self.build_command(root, {"id": worker_id, "port": 8190})
+                path = Path(cmd[cmd.index("--database-url") + 1].removeprefix("sqlite:///"))
+                self.assertTrue(path.is_relative_to(root / "user"))
+                databases.append(path)
+            self.assertEqual(len(set(databases)), 3)
+
+    def test_extra_args_cannot_silently_override_the_configured_port(self):
+        with tempfile.TemporaryDirectory() as directory:
+            root = Path(directory)
+            (root / "main.py").touch()
+            for extra in ["--port 8189", "--port=8189"]:
+                with self.subTest(extra=extra), self.assertRaisesRegex(ValueError, "configured worker port"):
+                    self.build_command(root, {"id": "worker-a", "port": 8190, "extra_args": extra})
+
+    def test_inherits_runtime_layout_args_for_desktop(self):
+        builder = launch_builder_module.LaunchCommandBuilder()
+        runtime_args = Namespace(
+            listen="127.0.0.1",
+            base_directory="C:/Users/test/ComfyUI",
+            temp_directory=None,
+            input_directory="C:/Users/test/ComfyUI/input",
+            output_directory="C:/Users/test/ComfyUI/output",
+            user_directory="C:/Users/test/ComfyUI/user",
+            front_end_root="C:/Program Files/ComfyUI/web_custom_versions/desktop_app",
+            extra_model_paths_config=[["C:/Users/test/AppData/Roaming/ComfyUI/extra_models_config.yaml"]],
+            enable_manager=True,
+            disable_manager_ui=False,
+            enable_manager_legacy_ui=False,
+            windows_standalone_build=True,
+            log_stdout=True,
+            verbose="INFO",
+            enable_cors_header="*",
+        )
+        comfy_module = types.ModuleType("comfy")
+        comfy_cli_args = types.ModuleType("comfy.cli_args")
+        comfy_cli_args.args = runtime_args
+
+        worker_config = {
+            "port": 9001,
+            "extra_args": "--preview-method auto",
+        }
+
+        def fake_exists(path):
+            return path == "/desktop/ComfyUI/main.py"
+
+        with patch.dict(
+            sys.modules,
+            {"comfy": comfy_module, "comfy.cli_args": comfy_cli_args},
+            clear=False,
+        ), patch.object(launch_builder_module.os.path, "exists", side_effect=fake_exists):
+            cmd = builder.build_launch_command(worker_config, "/desktop/ComfyUI")
+
+        self.assertEqual(cmd[:2], ["/usr/bin/test-python", "/desktop/ComfyUI/main.py"])
+        self.assertIn("--listen", cmd)
+        self.assertIn("127.0.0.1", cmd)
+        self.assertIn("--base-directory", cmd)
+        self.assertIn("C:/Users/test/ComfyUI", cmd)
+        self.assertIn("--input-directory", cmd)
+        self.assertIn("--output-directory", cmd)
+        self.assertIn("--user-directory", cmd)
+        self.assertIn("--front-end-root", cmd)
+        self.assertIn("--extra-model-paths-config", cmd)
+        self.assertIn("C:/Users/test/AppData/Roaming/ComfyUI/extra_models_config.yaml", cmd)
+        self.assertIn("--enable-manager", cmd)
+        self.assertIn("--windows-standalone-build", cmd)
+        self.assertIn("--log-stdout", cmd)
+        self.assertIn("--disable-auto-launch", cmd)
+        self.assertIn("--enable-cors-header", cmd)
+        self.assertIn("*", cmd)
+        self.assertIn("--port", cmd)
+        self.assertIn("9001", cmd)
+        self.assertNotIn("--auto-launch", cmd)
+
+
+class ProcessLaunchPortTests(unittest.TestCase):
+    def test_launch_before_master_listens_uses_configured_master_port(self):
+        lifecycle_module = _load_process_module("lifecycle.py")
+        cli_args = types.ModuleType("comfy.cli_args")
+        cli_args.args = Namespace(port=8189)
+        for missing in [AttributeError("PromptServer has no port"), None]:
+            with self.subTest(missing=missing), tempfile.TemporaryDirectory() as directory:
+                manager = types.SimpleNamespace(
+                    find_comfy_root=lambda: directory,
+                    build_launch_command=lambda _worker, _root: ["python", "main.py"],
+                    processes={}, save_processes=lambda: None,
+                )
+                with patch.dict(sys.modules, {"comfy": types.ModuleType("comfy"), "comfy.cli_args": cli_args}), \
+                     patch.object(lifecycle_module, "get_server_port", side_effect=missing if isinstance(missing, Exception) else None, return_value=None), \
+                     patch.object(lifecycle_module, "validate_worker_port") as validate, \
+                     patch.object(lifecycle_module.subprocess, "Popen", return_value=types.SimpleNamespace(pid=1234)) as spawn:
+                    pid = lifecycle_module.ProcessLifecycle(manager).launch_worker({"id": "manual", "port": 8190})
+                validate.assert_called_once_with({"id": "manual", "port": 8190}, 8189, [])
+                spawn.assert_called_once()
+                self.assertEqual(pid, 1234)
+
+    def test_pre_listen_launch_still_rejects_the_configured_master_port(self):
+        lifecycle_module = _load_process_module("lifecycle.py")
+        cli_args = types.ModuleType("comfy.cli_args")
+        cli_args.args = Namespace(port=8189)
+        manager = types.SimpleNamespace(find_comfy_root=lambda: "/ComfyUI", processes={})
+        with patch.dict(sys.modules, {"comfy": types.ModuleType("comfy"), "comfy.cli_args": cli_args}), \
+             patch.object(lifecycle_module, "get_server_port", side_effect=AttributeError("port")), \
+             patch.object(lifecycle_module.subprocess, "Popen") as spawn:
+            with self.assertRaisesRegex(ValueError, "master port 8189"):
+                lifecycle_module.ProcessLifecycle(manager).launch_worker({"id": "manual", "port": 8189})
+        spawn.assert_not_called()
+
+    def test_conflict_is_rejected_before_building_or_spawning(self):
+        lifecycle_module = _load_process_module("lifecycle.py")
+        manager = types.SimpleNamespace(find_comfy_root=lambda: "/ComfyUI", processes={})
+        with patch.object(lifecycle_module.subprocess, "Popen") as spawn:
+            with self.assertRaisesRegex(ValueError, "master port 8189"):
+                lifecycle_module.ProcessLifecycle(manager).launch_worker({"id": "manual", "port": 8189})
+        spawn.assert_not_called()
+        self.assertEqual(manager.processes, {})
+
+    def test_available_port_reaches_existing_launch_path(self):
+        lifecycle_module = _load_process_module("lifecycle.py")
+        with tempfile.TemporaryDirectory() as directory:
+            manager = types.SimpleNamespace(
+                find_comfy_root=lambda: directory,
+                build_launch_command=lambda _worker, _root: ["python", "main.py", "--port", "8190"],
+                processes={}, save_processes=lambda: None,
+            )
+            with patch.object(lifecycle_module, "validate_worker_port") as validate, \
+                 patch.object(lifecycle_module.subprocess, "Popen", return_value=types.SimpleNamespace(pid=1234)) as spawn:
+                pid = lifecycle_module.ProcessLifecycle(manager).launch_worker({"id": "manual", "port": 8190})
+            validate.assert_called_once_with({"id": "manual", "port": 8190}, 8189, [])
+            spawn.assert_called_once()
+            self.assertEqual(pid, 1234)
+            self.assertEqual(manager.processes["manual"]["pid"], 1234)
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/unit/__init__.py

@@ -1 +0,0 @@
-"""Unit tests."""

tests/unit/nodes/__init__.py

@@ -1 +0,0 @@
-"""Unit tests."""

tests/unit/nodes/test_branch.py

@@ -1,80 +0,0 @@
-import importlib.util
-import sys
-import types
-import unittest
-from pathlib import Path
-
-
-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
-
-
-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_branch_module():
-    package_name = "dist_branch_testpkg"
-    _bootstrap_package(package_name)
-    _load_module(package_name, "nodes/utilities.py", "nodes.utilities")
-    return _load_module(package_name, "nodes/branch.py", "nodes.branch")
-
-
-branch_module = _load_branch_module()
-
-
-class DistributedBranchTests(unittest.TestCase):
-    def test_node_declares_ten_output_slots(self):
-        self.assertEqual(len(branch_module.DistributedBranch.RETURN_TYPES), 10)
-        self.assertEqual(len(branch_module.DistributedBranch.RETURN_NAMES), 10)
-
-    def test_no_distribution_outputs_all_branches(self):
-        node = branch_module.DistributedBranch()
-        token = {"k": "v"}
-
-        outputs = node.branch(token, num_branches=2, assigned_branch=-1)
-
-        self.assertEqual(len(outputs), 10)
-        self.assertTrue(all(output is token for output in outputs))
-
-    def test_assigned_branch_still_passes_input(self):
-        node = branch_module.DistributedBranch()
-        token = "payload"
-
-        outputs = node.branch(token, num_branches=3, is_worker=True, worker_id="worker-a", assigned_branch=1)
-
-        self.assertEqual(outputs[0], "payload")
-        self.assertEqual(outputs[1], "payload")
-        self.assertEqual(outputs[9], "payload")
-
-
-if __name__ == "__main__":
-    unittest.main()

tests/unit/nodes/test_branch_collector.py

@@ -1,282 +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_branch_collector_module():
-    package_name = "dist_join_testpkg"
-    _bootstrap_package(package_name)
-    _load_module(package_name, "nodes/utilities.py", "nodes.utilities")
-    _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/branch_collector.py", "nodes.branch_collector")
-
-
-branch_collector_module = _load_branch_collector_module()
-
-
-def _context(**overrides):
-    defaults = {
-        "multi_job_id": "",
-        "is_worker": False,
-        "master_url": "",
-        "enabled_worker_ids": "[]",
-        "worker_id": "",
-        "assigned_branch": -1,
-        "delegate_only": False,
-    }
-    defaults.update(overrides)
-    return branch_collector_module.BranchRunContext(**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 DistributedBranchCollectorTests(unittest.IsolatedAsyncioTestCase):
-    def setUp(self):
-        self.node = branch_collector_module.DistributedBranchCollector()
-
-    async def test_worker_mode_sends_result_with_assigned_branch_index(self):
-        fake_session = _FakeSession()
-
-        async def _get_session():
-            return fake_session
-
-        branch_collector_module.get_client_session = _get_session
-
-        tensor = torch.ones((1, 4, 4, 3), dtype=torch.float32)
-        outputs = await self.node.execute(
-            [None, tensor, None, None, None, None, None, None, None, None],
-            num_branches=2,
-            context=_context(
-                multi_job_id="join-job",
-                is_worker=True,
-                master_url="http://master.local:8188",
-                worker_id="worker-a",
-                assigned_branch=1,
-            ),
-        )
-
-        self.assertEqual(len(fake_session.calls), 1)
-        payload = fake_session.calls[0]["payload"]
-        self.assertEqual(payload["batch_idx"], 1)
-        self.assertTrue(payload["is_last"])
-        self.assertEqual(payload["worker_id"], "worker-a")
-        self.assertTrue(torch.equal(outputs[1], tensor))
-
-    async def test_master_mode_collects_worker_results_into_branch_slots(self):
-        fake_server = _FakePromptServer()
-        branch_collector_module._get_prompt_server_instance = lambda: fake_server
-
-        queue = asyncio.Queue()
-        worker_1 = torch.full((1, 4, 4, 3), 0.5, dtype=torch.float32)
-        worker_2 = torch.full((1, 4, 4, 3), 0.8, dtype=torch.float32)
-        await queue.put({"tensor": worker_1, "worker_id": "worker-a", "image_index": 1, "is_last": True})
-        await queue.put({"tensor": worker_2, "worker_id": "worker-b", "image_index": 2, "is_last": True})
-        fake_server.distributed_pending_jobs["join-job"] = queue
-
-        master_tensor = torch.zeros((1, 4, 4, 3), dtype=torch.float32)
-        outputs = await self.node.execute(
-            [master_tensor, None, None, None, None, None, None, None, None, None],
-            num_branches=3,
-            context=_context(
-                multi_job_id="join-job",
-                enabled_worker_ids=json.dumps(["worker-a", "worker-b"]),
-                assigned_branch=0,
-            ),
-        )
-
-        self.assertTrue(torch.equal(outputs[0], master_tensor))
-        self.assertTrue(torch.equal(outputs[1], worker_1))
-        self.assertTrue(torch.equal(outputs[2], worker_2))
-
-    async def test_unassigned_slots_remain_none(self):
-        fake_server = _FakePromptServer()
-        branch_collector_module._get_prompt_server_instance = lambda: fake_server
-
-        queue = asyncio.Queue()
-        worker_1 = torch.full((1, 4, 4, 3), 0.5, dtype=torch.float32)
-        await queue.put({"tensor": worker_1, "worker_id": "worker-a", "image_index": 1, "is_last": True})
-        fake_server.distributed_pending_jobs["join-job"] = queue
-
-        master_tensor = torch.zeros((1, 4, 4, 3), dtype=torch.float32)
-        outputs = await self.node.execute(
-            [master_tensor, None, None, None, None, None, None, None, None, None],
-            num_branches=2,
-            context=_context(
-                multi_job_id="join-job",
-                enabled_worker_ids=json.dumps(["worker-a"]),
-                assigned_branch=0,
-            ),
-        )
-
-        self.assertIsNone(outputs[2])
-        self.assertIsNone(outputs[9])
-
-    async def test_missing_worker_results_use_safe_fallback_for_expected_slots(self):
-        fake_server = _FakePromptServer()
-        branch_collector_module._get_prompt_server_instance = lambda: fake_server
-
-        queue = asyncio.Queue()
-        # No worker payloads enqueued; master should timeout and fill expected worker slots.
-        fake_server.distributed_pending_jobs["join-job"] = queue
-
-        master_tensor = torch.full((1, 4, 4, 3), 0.7, dtype=torch.float32)
-        outputs = await self.node.execute(
-            [master_tensor, None, None, None, None, None, None, None, None, None],
-            num_branches=3,
-            context=_context(
-                multi_job_id="join-job",
-                enabled_worker_ids=json.dumps(["worker-a", "worker-b"]),
-                assigned_branch=0,
-            ),
-        )
-
-        self.assertTrue(torch.equal(outputs[0], master_tensor))
-        self.assertIsInstance(outputs[1], torch.Tensor)
-        self.assertIsInstance(outputs[2], torch.Tensor)
-        self.assertEqual(float(outputs[1].abs().sum().item()), 0.0)
-        self.assertEqual(float(outputs[2].abs().sum().item()), 0.0)
-
-    async def test_assigned_slot_can_be_filled_from_single_mismatched_local_input(self):
-        fake_server = _FakePromptServer()
-        branch_collector_module._get_prompt_server_instance = lambda: fake_server
-
-        # No worker results expected for this check.
-        fake_server.distributed_pending_jobs["join-job"] = asyncio.Queue()
-
-        local_tensor = torch.full((1, 4, 4, 3), 0.33, dtype=torch.float32)
-        outputs = await self.node.execute(
-            # Local value arrives on slot 1 even though participant is assigned slot 0.
-            [None, local_tensor, None, None, None, None, None, None, None, None],
-            num_branches=3,
-            context=_context(
-                multi_job_id="join-job",
-                enabled_worker_ids=json.dumps([]),
-                assigned_branch=0,
-            ),
-        )
-
-        self.assertTrue(torch.equal(outputs[0], local_tensor))
-
-
-if __name__ == "__main__":
-    unittest.main()

tests/unit/nodes/test_collector.py

@@ -1,228 +0,0 @@
-import asyncio
-import importlib.util
-import sys
-import types
-import unittest
-from pathlib import Path
-
-import numpy as np
-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.load_config = lambda: {"workers": []}
-    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 pil_to_tensor(image):
-        array = np.array(image).astype(np.float32) / 255.0
-        return torch.from_numpy(array).unsqueeze(0)
-
-    image_module.ensure_contiguous = ensure_contiguous
-    image_module.tensor_to_pil = tensor_to_pil
-    image_module.pil_to_tensor = pil_to_tensor
-    sys.modules[f"{package_name}.utils.image"] = image_module
-
-    network_module = types.ModuleType(f"{package_name}.utils.network")
-    network_module.build_worker_url = lambda _worker: "http://worker.local:8188"
-    network_module.get_client_session = lambda: None
-
-    async def probe_worker(_url, timeout=2.0):
-        _ = timeout
-        return None
-
-    network_module.probe_worker = probe_worker
-    sys.modules[f"{package_name}.utils.network"] = network_module
-
-    audio_payload_module = types.ModuleType(f"{package_name}.utils.audio_payload")
-    audio_payload_module.encode_audio_payload = lambda payload: payload
-    sys.modules[f"{package_name}.utils.audio_payload"] = audio_payload_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
-
-    if "aiohttp" not in sys.modules:
-        try:
-            import aiohttp as _aiohttp  # noqa: F401
-        except Exception:
-            aiohttp_module = types.ModuleType("aiohttp")
-
-            class _ClientTimeout:
-                def __init__(self, total=None):
-                    self.total = total
-
-            aiohttp_module.ClientTimeout = _ClientTimeout
-            sys.modules["aiohttp"] = aiohttp_module
-
-    class _FakePromptServer:
-        def __init__(self):
-            self.distributed_pending_jobs = {}
-            self.distributed_jobs_lock = asyncio.Lock()
-
-    server_module = types.ModuleType("server")
-    server_module.PromptServer = types.SimpleNamespace(instance=_FakePromptServer())
-    sys.modules["server"] = server_module
-
-    comfy_module = types.ModuleType("comfy")
-    model_mgmt = types.ModuleType("comfy.model_management")
-
-    class _InterruptProcessingException(Exception):
-        pass
-
-    model_mgmt.throw_exception_if_processing_interrupted = lambda: None
-    model_mgmt.InterruptProcessingException = _InterruptProcessingException
-    comfy_module.model_management = model_mgmt
-    sys.modules["comfy"] = comfy_module
-    sys.modules["comfy.model_management"] = model_mgmt
-
-    comfy_utils_module = types.ModuleType("comfy.utils")
-
-    class _ProgressBar:
-        def __init__(self, _total):
-            self.total = _total
-
-        def update(self, _step):
-            return None
-
-    comfy_utils_module.ProgressBar = _ProgressBar
-    sys.modules["comfy.utils"] = comfy_utils_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_collector_module():
-    package_name = "dist_collector_testpkg"
-    _bootstrap_package(package_name)
-    _load_module(package_name, "nodes/hidden_inputs.py", "nodes.hidden_inputs")
-    _load_module(package_name, "utils/worker_ids.py", "utils.worker_ids")
-    return _load_module(package_name, "nodes/collector.py", "nodes.collector")
-
-
-collector_module = _load_collector_module()
-
-
-class DistributedCollectorTests(unittest.TestCase):
-    def setUp(self):
-        self.node = collector_module.DistributedCollectorNode()
-
-    def test_store_worker_result_tracks_by_worker_and_index(self):
-        worker_images = {}
-        stored = self.node._store_worker_result(
-            worker_images,
-            {
-                "worker_id": "worker-a",
-                "image_index": 2,
-                "tensor": torch.full((1, 1, 1, 1), 0.25),
-            },
-        )
-
-        self.assertEqual(stored, 1)
-        self.assertIn("worker-a", worker_images)
-        self.assertIn(2, worker_images["worker-a"])
-
-    def test_combine_audio_honors_master_then_worker_order(self):
-        master_audio = {"waveform": torch.ones((1, 2, 1)), "sample_rate": 48000}
-        worker_audio = {
-            "worker-b": {"waveform": torch.full((1, 2, 1), 3.0), "sample_rate": 48000},
-            "worker-a": {"waveform": torch.full((1, 2, 2), 2.0), "sample_rate": 48000},
-        }
-        empty_audio = {"waveform": torch.zeros((1, 2, 1)), "sample_rate": 44100}
-
-        combined = self.node._combine_audio(
-            master_audio=master_audio,
-            worker_audio=worker_audio,
-            empty_audio=empty_audio,
-            worker_order=["worker-a", "worker-b"],
-        )
-
-        expected = torch.cat(
-            [
-                master_audio["waveform"],
-                worker_audio["worker-a"]["waveform"],
-                worker_audio["worker-b"]["waveform"],
-            ],
-            dim=-1,
-        )
-        self.assertEqual(combined["sample_rate"], 48000)
-        self.assertTrue(torch.equal(combined["waveform"], expected))
-
-    def test_reorder_and_combine_tensors_uses_enabled_worker_priority(self):
-        worker_images = {
-            "worker-b": {
-                1: torch.full((1, 1, 1, 1), 4.0),
-                0: torch.full((1, 1, 1, 1), 3.0),
-            },
-            "worker-a": {
-                0: torch.full((1, 1, 1, 1), 2.0),
-            },
-        }
-        master_images = torch.full((1, 1, 1, 1), 1.0)
-
-        combined = self.node._reorder_and_combine_tensors(
-            worker_images=worker_images,
-            worker_order=["worker-a", "worker-b"],
-            master_batch_size=1,
-            images_on_cpu=master_images,
-            delegate_mode=False,
-            fallback_images=None,
-        )
-
-        self.assertEqual(combined.shape, (4, 1, 1, 1))
-        self.assertEqual(combined[0, 0, 0, 0].item(), 1.0)
-        self.assertEqual(combined[1, 0, 0, 0].item(), 2.0)
-        self.assertEqual(combined[2, 0, 0, 0].item(), 3.0)
-        self.assertEqual(combined[3, 0, 0, 0].item(), 4.0)
-
-
-if __name__ == "__main__":
-    unittest.main()

tests/unit/nodes/test_distributed_upscale.py

@@ -1,248 +0,0 @@
-import importlib.util
-import sys
-import types
-import unittest
-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
-
-    modes_pkg = types.ModuleType(f"{package_name}.upscale.modes")
-    modes_pkg.__path__ = []
-    sys.modules[f"{package_name}.upscale.modes"] = modes_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
-
-    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})")
-    )
-    sys.modules[f"{package_name}.utils.async_helpers"] = async_helpers_module
-
-    job_store_module = types.ModuleType(f"{package_name}.upscale.job_store")
-    job_store_module.ensure_tile_jobs_initialized = lambda: None
-    sys.modules[f"{package_name}.upscale.job_store"] = job_store_module
-
-    tile_ops_module = types.ModuleType(f"{package_name}.upscale.tile_ops")
-
-    class _TileOpsMixin:
-        def round_to_multiple(self, value):
-            return int(value)
-
-        def calculate_tiles(self, width, height, tile_width, tile_height, force_uniform_tiles):
-            _ = force_uniform_tiles
-            tiles = []
-            for y in range(0, int(height), max(1, int(tile_height))):
-                for x in range(0, int(width), max(1, int(tile_width))):
-                    tiles.append((x, y))
-            return tiles or [(0, 0)]
-
-    tile_ops_module.TileOpsMixin = _TileOpsMixin
-    sys.modules[f"{package_name}.upscale.tile_ops"] = tile_ops_module
-
-    result_collector_module = types.ModuleType(f"{package_name}.upscale.result_collector")
-    result_collector_module.ResultCollectorMixin = type("ResultCollectorMixin", (), {})
-    sys.modules[f"{package_name}.upscale.result_collector"] = result_collector_module
-
-    worker_comms_module = types.ModuleType(f"{package_name}.upscale.worker_comms")
-    worker_comms_module.WorkerCommsMixin = type("WorkerCommsMixin", (), {})
-    sys.modules[f"{package_name}.upscale.worker_comms"] = worker_comms_module
-
-    job_state_module = types.ModuleType(f"{package_name}.upscale.job_state")
-    job_state_module.JobStateMixin = type("JobStateMixin", (), {})
-    sys.modules[f"{package_name}.upscale.job_state"] = job_state_module
-
-    single_gpu_module = types.ModuleType(f"{package_name}.upscale.modes.single_gpu")
-
-    class _SingleGpuModeMixin:
-        def process_single_gpu(self, *_args, **_kwargs):
-            return ("single_gpu_result",)
-
-    single_gpu_module.SingleGpuModeMixin = _SingleGpuModeMixin
-    sys.modules[f"{package_name}.upscale.modes.single_gpu"] = single_gpu_module
-
-    static_mode_module = types.ModuleType(f"{package_name}.upscale.modes.static")
-
-    class _StaticModeMixin:
-        def _process_worker_static_sync(self, *_args, **_kwargs):
-            return ("worker_static_result",)
-
-        def _process_master_static_sync(self, *_args, **_kwargs):
-            return ("master_static_result",)
-
-    static_mode_module.StaticModeMixin = _StaticModeMixin
-    sys.modules[f"{package_name}.upscale.modes.static"] = static_mode_module
-
-    dynamic_mode_module = types.ModuleType(f"{package_name}.upscale.modes.dynamic")
-
-    class _DynamicModeMixin:
-        def process_worker_dynamic(self, *_args, **_kwargs):
-            return ("worker_dynamic_result",)
-
-        def process_master_dynamic(self, *_args, **_kwargs):
-            return ("master_dynamic_result",)
-
-    dynamic_mode_module.DynamicModeMixin = _DynamicModeMixin
-    sys.modules[f"{package_name}.upscale.modes.dynamic"] = dynamic_mode_module
-
-    comfy_module = types.ModuleType("comfy")
-    samplers_module = types.ModuleType("comfy.samplers")
-
-    class _KSampler:
-        SAMPLERS = ("euler",)
-        SCHEDULERS = ("normal",)
-
-    samplers_module.KSampler = _KSampler
-    comfy_module.samplers = samplers_module
-    sys.modules["comfy"] = comfy_module
-    sys.modules["comfy.samplers"] = samplers_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_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")
-
-
-upscale_module = _load_distributed_upscale_module()
-
-
-class DistributedUpscaleTests(unittest.TestCase):
-    def setUp(self):
-        self.node = upscale_module.UltimateSDUpscaleDistributed()
-        self.base_image = torch.zeros((1, 64, 64, 3), dtype=torch.float32)
-        self.common_kwargs = {
-            "model": object(),
-            "positive": object(),
-            "negative": object(),
-            "vae": object(),
-            "seed": 1,
-            "steps": 10,
-            "cfg": 7.5,
-            "sampler_name": "euler",
-            "scheduler": "normal",
-            "denoise": 0.5,
-            "tile_width": 32,
-            "tile_height": 32,
-            "padding": 16,
-            "mask_blur": 4,
-            "force_uniform_tiles": True,
-            "tiled_decode": False,
-        }
-
-    def test_parse_enabled_worker_ids_supports_json_and_lists(self):
-        self.assertEqual(
-            upscale_module._parse_enabled_worker_ids(["worker-a", 2]),
-            ["worker-a", "2"],
-        )
-        self.assertEqual(
-            upscale_module._parse_enabled_worker_ids('["worker-a","worker-b"]'),
-            ["worker-a", "worker-b"],
-        )
-        self.assertEqual(upscale_module._parse_enabled_worker_ids("invalid-json"), [])
-        self.assertEqual(upscale_module._parse_enabled_worker_ids(None), [])
-
-    def test_determine_processing_mode(self):
-        self.assertEqual(self.node._determine_processing_mode(batch_size=1, num_workers=0, dynamic_threshold=8), "single_gpu")
-        self.assertEqual(self.node._determine_processing_mode(batch_size=9, num_workers=2, dynamic_threshold=8), "static")
-
-    def test_run_raises_for_non_4n_plus_1_master_batches(self):
-        bad_batch = torch.zeros((2, 64, 64, 3), dtype=torch.float32)
-        with self.assertRaises(ValueError):
-            self.node.run(
-                upscaled_image=bad_batch,
-                multi_job_id="job-1",
-                is_worker=False,
-                **self.common_kwargs,
-            )
-
-    def test_run_dispatches_single_gpu_when_not_distributed(self):
-        self.node.process_single_gpu = lambda *_args, **_kwargs: ("single",)
-        result = self.node.run(
-            upscaled_image=self.base_image,
-            multi_job_id="",
-            is_worker=False,
-            **self.common_kwargs,
-        )
-        self.assertEqual(result, ("single",))
-
-    def test_run_dispatches_worker_or_master_paths(self):
-        self.node.process_worker = lambda *_args, **_kwargs: ("worker",)
-        self.node.process_master = lambda *_args, **_kwargs: ("master",)
-
-        worker_result = self.node.run(
-            upscaled_image=self.base_image,
-            multi_job_id="job-2",
-            is_worker=True,
-            master_url="http://master.local:8188",
-            enabled_worker_ids='["worker-a"]',
-            worker_id="worker-a",
-            tile_indices="",
-            dynamic_threshold=8,
-            **self.common_kwargs,
-        )
-        master_result = self.node.run(
-            upscaled_image=self.base_image,
-            multi_job_id="job-2",
-            is_worker=False,
-            enabled_worker_ids='["worker-a"]',
-            worker_id="",
-            tile_indices="",
-            dynamic_threshold=8,
-            **self.common_kwargs,
-        )
-
-        self.assertEqual(worker_result, ("worker",))
-        self.assertEqual(master_result, ("master",))
-
-
-if __name__ == "__main__":
-    unittest.main()

tests/unit/system/__init__.py

@@ -1 +0,0 @@
-"""Unit tests."""

tests/unit/system/test_distributed.py

@@ -1,71 +0,0 @@
-import importlib.util
-import sys
-import types
-import unittest
-from pathlib import Path
-
-
-def _load_distributed_module():
-    module_path = Path(__file__).resolve().parents[3] / "distributed.py"
-    package_name = "dist_distributed_entry_testpkg"
-    calls = {
-        "build_node_mappings": 0,
-        "initialize_runtime": 0,
-        "initialize_runtime_args": [],
-    }
-
-    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
-
-    bootstrap_pkg = types.ModuleType(f"{package_name}.bootstrap")
-    bootstrap_pkg.__path__ = []
-    sys.modules[f"{package_name}.bootstrap"] = bootstrap_pkg
-
-    entrypoint_module = types.ModuleType(f"{package_name}.bootstrap.entrypoint")
-
-    def _build_node_mappings():
-        calls["build_node_mappings"] += 1
-        return ({"MockNode": object}, {"MockNode": "Mock Node"})
-
-    def _initialize_runtime(prompt_server=None):
-        calls["initialize_runtime"] += 1
-        calls["initialize_runtime_args"].append(prompt_server)
-
-    entrypoint_module.build_node_mappings = _build_node_mappings
-    entrypoint_module.initialize_runtime = _initialize_runtime
-    sys.modules[f"{package_name}.bootstrap.entrypoint"] = entrypoint_module
-
-    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_builds_mappings_without_runtime_bootstrap(self):
-        module, calls = _load_distributed_module()
-
-        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_initialize_runtime_delegates_to_bootstrap_entrypoint(self):
-        module, calls = _load_distributed_module()
-        sentinel = object()
-
-        module.initialize_runtime(sentinel)
-
-        self.assertEqual(calls["initialize_runtime"], 1)
-        self.assertEqual(calls["initialize_runtime_args"], [sentinel])
-
-
-if __name__ == "__main__":
-    unittest.main()