fix: preserve image alpha in distributed collection

README.md

@@ -27,11 +27,6 @@
 - Intelligent distribution
 - Handles single images and videos
 
-#### Dynamic Image-List Processing
-- Process a whole image list in parallel: the master and every worker take the next image when they free up, and results come back in input order
-- Use **Distributed Image List Start / End** around a per-image processing region
-- [Wiring, supported inputs, and first-version limits](/docs/image-list-queue.md)
-
 #### Ease of Use
 - Auto-setup local workers; easily add remote/cloud ones
 - Convert any workflow to distributed with 2 nodes
@@ -43,12 +38,11 @@
 
 <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 four types of workers:
+ComfyUI Distributed supports three 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 endpoints for image Collector jobs; no Distributed install or inbound tunnel on the endpoint. [Setup, endpoint contract, and preview limitations](/docs/serverless-workers.md).
 
 > For detailed setup instructions, see the [setup guide](/docs/worker-setup-guides.md)
 
@@ -184,8 +178,6 @@ Use **Distributed Value** when you want per-worker overrides (for example, diffe
 | **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 |
-| **Distributed Image List Start** | Turns an image list or batch into one job per image, scheduled to the next available participant; mask and value inputs are shared across images or paired one per image |
-| **Distributed Image List End** | Closes the paired region and returns images, masks, and texts as lists in the original input order |
 
 ---
 
@@ -228,26 +220,6 @@ 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.

__init__.py

@@ -1,22 +1,29 @@
-"""ComfyUI-Distributed's native V3 extension entrypoint."""
-from comfy_api.v0_0_2 import ComfyExtension, io
+# Import everything needed from the main module
+from .distributed import (
+    NODE_CLASS_MAPPINGS as DISTRIBUTED_CLASS_MAPPINGS, 
+    NODE_DISPLAY_NAME_MAPPINGS as DISTRIBUTED_DISPLAY_NAME_MAPPINGS
+)
 
-from .nodes.v3 import NODES
-from .runtime.bootstrap import initialize
+# Import utilities
+from .utils.config import ensure_config_exists, CONFIG_FILE
+from .utils.logging import debug_log
 
-WEB_DIRECTORY = './web'
+# Import distributed upscale nodes
+from .nodes.distributed_upscale import (
+    NODE_CLASS_MAPPINGS as UPSCALE_CLASS_MAPPINGS,
+    NODE_DISPLAY_NAME_MAPPINGS as UPSCALE_DISPLAY_NAME_MAPPINGS
+)
 
+WEB_DIRECTORY = "./web"
 
-class DistributedExtension(ComfyExtension):
-    async def on_load(self) -> None:
-        initialize()
+ensure_config_exists()
 
-    async def get_node_list(self) -> list[type[io.ComfyNode]]:
-        return list(NODES)
+# Merge node mappings
+NODE_CLASS_MAPPINGS = {**DISTRIBUTED_CLASS_MAPPINGS, **UPSCALE_CLASS_MAPPINGS}
+NODE_DISPLAY_NAME_MAPPINGS = {**DISTRIBUTED_DISPLAY_NAME_MAPPINGS, **UPSCALE_DISPLAY_NAME_MAPPINGS}
 
+__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS']
 
-async def comfy_entrypoint() -> DistributedExtension:
-    return DistributedExtension()
-
-
-__all__ = ['comfy_entrypoint', 'WEB_DIRECTORY']
+debug_log("Loaded Distributed nodes.")
+debug_log(f"Config file: {CONFIG_FILE}")
+debug_log(f"Available nodes: {list(NODE_CLASS_MAPPINGS.keys())}")

api/__init__.py

@@ -1,7 +1,5 @@
 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
-from . import image_list_routes  # noqa: F401

api/config_routes.py

@@ -1,6 +1,5 @@
 import json
 from contextlib import asynccontextmanager
-from typing import Any
 
 from aiohttp import web
 import server
@@ -25,12 +24,6 @@ except ImportError:
             save_config(config)
 from ..utils.logging import debug_log, log
 from ..utils.network import handle_api_error, normalize_host
-from ..utils.secret_fields import is_secret_field_name
-from ..utils.serverless_config import (
-    ServerlessSecretChanges,
-    public_config,
-    validate_serverless_worker,
-)
 
 
 def _positive_int(value):
@@ -93,159 +86,22 @@ 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
 
 
-_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",
-}
-
-
-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:
-            validated.append(worker)
-    return validated
-
-
-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):
-    try:
-        config = load_config()
-        return web.json_response(public_config(config, _serverless_secret_store()))
-    except Exception:
-        return _error_response("Unable to read config", 500)
+    config = load_config()
+    return web.json_response(config)
 
 
 @server.PromptServer.instance.routes.post("/distributed/config")
 async def update_config_endpoint(request):
-    """Bulk config update with schema validation and credential redaction."""
+    """Bulk config update with schema validation."""
     try:
         data = await request.json()
-    except Exception:
-        return _error_response("Invalid JSON payload", 400)
+    except Exception as e:
+        return await handle_api_error(request, f"Invalid JSON payload: {e}", 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)
+        return await handle_api_error(request, "Config payload must be an object", 400)
 
     validated_settings = {}
     validated_root = {}
@@ -262,15 +118,9 @@ async def update_config_endpoint(request):
             continue
 
         if validator and not validator(value):
-            errors.append(f"{key}: failed validation")
+            errors.append(f"{key}: value {value!r} failed validation")
             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:
@@ -283,41 +133,15 @@ async def update_config_endpoint(request):
             "message": "; ".join(errors),
         }, status=400)
 
-    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)
-
-                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)
+        async with config_transaction() as config:
+            settings = config.setdefault("settings", {})
+            settings.update(validated_settings)
+            for key, value in validated_root.items():
+                config[key] = value
+            return web.json_response({"status": "success", "config": config})
+    except Exception as e:
+        return await handle_api_error(request, e)
 
 
 @server.PromptServer.instance.routes.get("/distributed/queue_status/{job_id}")
@@ -340,202 +164,79 @@ async def queue_status_endpoint(request):
 
 @server.PromptServer.instance.routes.post("/distributed/config/update_worker")
 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)
-
+        data = await request.json()
         worker_id = data.get("worker_id")
+        
         if worker_id is None:
-            return _error_response("Missing worker_id", 400)
+            return await handle_api_error(request, "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")
 
-            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)
+            for worker in workers:
+                if worker["id"] == worker_id:
+                    _apply_field_patch(worker, data, _WORKER_FIELDS)
+                    worker_found = True
+                    break
 
-                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)
+            if not worker_found:
+                # If worker not found and all required fields are provided, create new worker
+                if all(key in data for key in ["name", "port", "cuda_device"]):
+                    new_worker = {
+                        "id": worker_id,
+                        "name": data["name"],
+                        "host": normalize_host(data.get("host", "localhost")),
+                        "port": data["port"],
+                        "cuda_device": data["cuda_device"],
+                        "enabled": data.get("enabled", False),
+                        "extra_args": data.get("extra_args", ""),
+                        "type": data.get("type", "local")
+                    }
+                    workers.append(new_worker)
                 else:
-                    existing.clear()
-                    existing.update(validated)
-                result = {"status": "success", "api_key_configured": next_secret is not None}
-
-            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 handle_api_error(
+                        request,
+                        f"Worker {worker_id} not found and missing required fields for creation",
+                        404,
+                    )
 
+            return web.json_response({"status": "success"})
+    except Exception as e:
+        return await handle_api_error(request, e, 400)
 
 @server.PromptServer.instance.routes.post("/distributed/config/delete_worker")
 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)
+        data = await request.json()
         worker_id = data.get("worker_id")
+        
         if worker_id is None:
-            return _error_response("Missing worker_id", 400)
-
-        result = None
+            return await handle_api_error(request, "Missing worker_id", 400)
+            
         async with config_transaction() as config:
             workers = config.get("workers", [])
-            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": "Worker deleted" if _worker_type(removed) == "serverless"
-                else f"Worker {removed.get('name', worker_id)} deleted",
-            }
 
-        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)
+            # Find and remove the worker
+            worker_index = -1
+            for i, worker in enumerate(workers):
+                if worker["id"] == worker_id:
+                    worker_index = i
+                    break
+
+            if worker_index == -1:
+                return await handle_api_error(request, f"Worker {worker_id} not found", 404)
+
+            # Remove the worker
+            removed_worker = workers.pop(worker_index)
+
+            return web.json_response({
+                "status": "success",
+                "message": f"Worker {removed_worker.get('name', worker_id)} deleted"
+            })
+    except Exception as e:
+        return await handle_api_error(request, e, 400)
 
 @server.PromptServer.instance.routes.post("/distributed/config/update_setting")
 async def update_setting_endpoint(request):

api/image_list_routes.py

@@ -1,89 +0,0 @@
-"""Bounded, typed HTTP transport for dynamic image-list worker tasks."""
-import json
-
-from aiohttp import web
-import server
-
-from ..image_list.runtime import (
-    MAX_REQUEST_BYTES,
-    PROTOCOL_VERSION,
-    decode_result,
-    delete_worker_item,
-    set_master_result,
-    store_worker_item,
-)
-
-prompt_server = server.PromptServer.instance
-
-
-async def _json_body(request):
-    length = request.content_length
-    if length is not None and length > MAX_REQUEST_BYTES:
-        raise ValueError(f"Request exceeds the {MAX_REQUEST_BYTES}-byte image-list limit")
-    raw = bytearray()
-    async for chunk in request.content.iter_chunked(64 * 1024):
-        raw.extend(chunk)
-        if len(raw) > MAX_REQUEST_BYTES:
-            raise ValueError(f"Request exceeds the {MAX_REQUEST_BYTES}-byte image-list limit")
-    try:
-        return json.loads(raw.decode("utf-8"))
-    except (UnicodeDecodeError, json.JSONDecodeError) as exc:
-        raise ValueError(f"Invalid image-list JSON body: {exc}") from exc
-
-
-@prompt_server.routes.get("/distributed/image-list/capabilities")
-async def image_list_capabilities_endpoint(_request):
-    return web.json_response({
-        "supported": True,
-        "protocol_version": PROTOCOL_VERSION,
-        "node_ids": ["DistributedImageListStart", "DistributedImageListEnd"],
-        "transport": "typed-float-tensor-json-v1",
-    })
-
-
-@prompt_server.routes.post("/distributed/image-list/item")
-async def image_list_item_endpoint(request):
-    try:
-        data = await _json_body(request)
-        if not isinstance(data, dict):
-            raise ValueError("Expected an image-list item object")
-        context_id = data.get("context_id")
-        if not isinstance(context_id, str) or not context_id:
-            raise ValueError("context_id must be a non-empty string")
-        item = data.get("item")
-        store_worker_item(data.get("token"), item)
-        return web.json_response({"status": "stored", "item_index": item["index"]})
-    except (ValueError, TypeError) as exc:
-        return web.json_response({"error": str(exc)}, status=400)
-    except Exception as exc:
-        return web.json_response({"error": str(exc)}, status=500)
-
-
-@prompt_server.routes.delete("/distributed/image-list/item/{token}")
-async def image_list_delete_item_endpoint(request):
-    token = request.match_info.get("token", "")
-    return web.json_response({"status": "deleted", "removed": delete_worker_item(token)})
-
-
-@prompt_server.routes.post("/distributed/image-list/result")
-async def image_list_result_endpoint(request):
-    try:
-        data = await _json_body(request)
-        if not isinstance(data, dict):
-            raise ValueError("Expected an image-list result object")
-        context_id = data.get("context_id")
-        token = data.get("token")
-        item_index = data.get("item_index")
-        if not isinstance(context_id, str) or not context_id:
-            raise ValueError("context_id must be a non-empty string")
-        if not isinstance(token, str) or not token:
-            raise ValueError("token must be a non-empty string")
-        if isinstance(item_index, bool) or not isinstance(item_index, int) or item_index < 0:
-            raise ValueError("item_index must be a non-negative integer")
-        result = decode_result(data.get("result"))
-        set_master_result(context_id, item_index, token, result)
-        return web.json_response({"status": "accepted", "item_index": item_index})
-    except (ValueError, TypeError) as exc:
-        return web.json_response({"error": str(exc)}, status=400)
-    except Exception as exc:
-        return web.json_response({"error": str(exc)}, status=500)

api/job_routes.py

@@ -15,7 +15,6 @@ 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 import transfer_codec
 try:
     from .queue_orchestration import ensure_distributed_state, orchestrate_distributed_execution
 except ImportError:
@@ -27,16 +26,8 @@ from .queue_request import parse_queue_request_payload
 
 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
+# Canonical worker result envelope accepted by POST /distributed/job_complete:
+# { "job_id": str, "worker_id": str, "batch_idx": int, "image": <base64 PNG>, "is_last": bool }
 
 
 def _decode_image_sync(image_path):
@@ -134,7 +125,9 @@ def _decode_canonical_png_tensor(image_payload):
 
     try:
         with Image.open(io.BytesIO(png_bytes)) as img:
-            img = img.convert("RGB")
+            # PNG callbacks must retain alpha just like local IMAGE tensors.
+            has_alpha = "A" in img.getbands() or "transparency" in img.info
+            img = img.convert("RGBA" if has_alpha else "RGB")
             tensor = pil_to_tensor(img)
         return ensure_contiguous(tensor)
     except Exception as exc:
@@ -279,136 +272,8 @@ async def check_file_endpoint(request):
         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):
-    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:
@@ -453,20 +318,28 @@ async def job_complete_endpoint(request):
         multi_job_id = job_id.strip()
         worker_id = worker_id.strip()
 
-        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)
+        pending = None
+        queue_size = 0
+        deadline = time.monotonic() + float(JOB_INIT_GRACE_PERIOD)
+        while pending is None:
+            async with prompt_server.distributed_jobs_lock:
+                pending = prompt_server.distributed_pending_jobs.get(multi_job_id)
+                if pending is not None:
+                    await pending.put(
+                        {
+                            "tensor": tensor,
+                            "worker_id": worker_id,
+                            "image_index": int(batch_idx),
+                            "is_last": is_last,
+                            "audio": decoded_audio,
+                        }
+                    )
+                    queue_size = pending.qsize()
+                    break
+
+            if time.monotonic() > deadline:
+                return await handle_api_error(request, "job not initialized", 404)
+            await asyncio.sleep(0.05)
 
         debug_log(
             f"job_complete received canonical envelope - job_id: {multi_job_id}, "

api/orchestration/dispatch.py

@@ -6,9 +6,6 @@ import aiohttp
 
 from ...utils.logging import debug_log, log
 from ...utils.network import build_worker_url, get_client_session, probe_worker
-# Must match image_list.runtime.PROTOCOL_VERSION (kept local so this module
-# stays importable without the image-list runtime; a test pins them equal).
-IMAGE_LIST_PROTOCOL_VERSION = 2
 try:
     from ...utils.trace_logger import trace_debug, trace_info
 except ImportError:
@@ -58,11 +55,6 @@ async def worker_ws_is_active(worker):
 
 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
 
@@ -80,44 +72,27 @@ async def _dispatch_via_websocket(worker_url, payload, client_id, timeout=60.0):
     ws_url = worker_url.replace("http://", "ws://").replace("https://", "wss://")
     ws_url = f"{ws_url}/distributed/worker_ws"
     session = await get_client_session()
-    loop = asyncio.get_running_loop()
-    deadline = loop.time() + timeout
-    ws = None
 
-    try:
-        async with asyncio.timeout_at(deadline):
-            ws = await session.ws_connect(ws_url, heartbeat=20)
-            await ws.send_json(ws_payload)
-            while True:
-                msg = await ws.receive()
-                if msg.type == aiohttp.WSMsgType.TEXT:
-                    data = json.loads(msg.data or "{}")
-                    if data.get("type") == "dispatch_ack" and data.get("request_id") == request_id:
-                        if data.get("ok"):
-                            return data.get("prompt_id")
-                        error_text = data.get("error") or "Worker rejected websocket dispatch."
-                        validation_error = data.get("validation_error")
-                        node_errors = data.get("node_errors")
-                        if validation_error:
-                            error_text = f"{error_text} | validation_error={validation_error}"
-                        if node_errors:
-                            error_text = f"{error_text} | node_errors={node_errors}"
-                        raise RuntimeError(error_text)
-                elif msg.type in (aiohttp.WSMsgType.ERROR, aiohttp.WSMsgType.CLOSE, aiohttp.WSMsgType.CLOSED):
-                    raise RuntimeError(f"Worker websocket closed before acknowledgement: {msg.type}; submission status is ambiguous")
-    except asyncio.TimeoutError as exc:
-        raise RuntimeError(
-            "Worker websocket dispatch timed out before dispatch_ack; submission status is ambiguous."
-        ) from exc
-    finally:
-        if ws is not None:
-            # Closing must not discard an already acknowledged prompt ID or
-            # replace a definite rejection with an ambiguous timeout.
-            try:
-                async with asyncio.timeout_at(min(deadline, loop.time() + 1.0)):
-                    await ws.close()
-            except (asyncio.TimeoutError, aiohttp.ClientError):
-                pass
+    async with session.ws_connect(ws_url, heartbeat=20, timeout=timeout) as ws:
+        await ws.send_json(ws_payload)
+        async for msg in ws:
+            if msg.type == aiohttp.WSMsgType.TEXT:
+                data = json.loads(msg.data or "{}")
+                if data.get("type") == "dispatch_ack" and data.get("request_id") == request_id:
+                    if data.get("ok"):
+                        return
+                    error_text = data.get("error") or "Worker rejected websocket dispatch."
+                    validation_error = data.get("validation_error")
+                    node_errors = data.get("node_errors")
+                    if validation_error:
+                        error_text = f"{error_text} | validation_error={validation_error}"
+                    if node_errors:
+                        error_text = f"{error_text} | node_errors={node_errors}"
+                    raise RuntimeError(error_text)
+            elif msg.type in (aiohttp.WSMsgType.ERROR, aiohttp.WSMsgType.CLOSED):
+                raise RuntimeError(f"Worker websocket closed unexpectedly: {msg.type}")
+
+    raise RuntimeError("Worker websocket closed before dispatch_ack was received.")
 
 
 async def dispatch_worker_prompt(
@@ -127,7 +102,6 @@ async def dispatch_worker_prompt(
     client_id=None,
     use_websocket=False,
     trace_execution_id=None,
-    return_prompt_id=False,
 ):
     """Send the prepared prompt to a worker ComfyUI instance."""
     worker_url = build_worker_url(worker)
@@ -141,7 +115,7 @@ async def dispatch_worker_prompt(
 
     if use_websocket:
         try:
-            prompt_id = await _dispatch_via_websocket(
+            await _dispatch_via_websocket(
                 worker_url,
                 {
                     "prompt": prompt_obj,
@@ -149,9 +123,7 @@ async def dispatch_worker_prompt(
                 },
                 client_id,
             )
-            if return_prompt_id and (not isinstance(prompt_id, str) or not prompt_id):
-                raise RuntimeError("Worker websocket acknowledgement omitted prompt_id")
-            return prompt_id if return_prompt_id else None
+            return
         except Exception as exc:
             worker_id = worker.get("id")
             if trace_execution_id:
@@ -167,96 +139,6 @@ async def dispatch_worker_prompt(
         timeout=aiohttp.ClientTimeout(total=60),
     ) as resp:
         resp.raise_for_status()
-        if not return_prompt_id:
-            return None
-        try:
-            response_data = await resp.json()
-        except Exception as exc:
-            raise RuntimeError("Worker /prompt response did not contain JSON acknowledgement") from exc
-        prompt_id = response_data.get("prompt_id") if isinstance(response_data, dict) else None
-        if not isinstance(prompt_id, str) or not prompt_id:
-            raise RuntimeError("Worker /prompt response omitted prompt_id")
-        return prompt_id if return_prompt_id else None
-
-
-async def preflight_image_list_workers(workers, require_idle=True, timeout=3.0):
-    """Require protocol/node support and report which direct workers are idle.
-
-    This performs read-only GETs only; it never submits a prompt or contacts a
-    serverless provider. Unsupported worker releases fail before master inputs run.
-    """
-    if not workers:
-        raise RuntimeError("Distributed Image List requires at least one selected direct worker")
-    session = await get_client_session()
-    required_nodes = {"DistributedImageListStart", "DistributedImageListEnd"}
-    semaphore = asyncio.Semaphore(8)
-
-    async def inspect(worker):
-        worker_id = str(worker.get("id") or "unknown")
-        check = "capability"
-        try:
-            async with semaphore:
-                cap_url = build_worker_url(worker, "/distributed/image-list/capabilities")
-                async with session.get(cap_url, timeout=aiohttp.ClientTimeout(total=timeout)) as response:
-                    if response.status != 200:
-                        raise RuntimeError(
-                            f"Image-list worker {worker_id} does not expose the capability endpoint "
-                            f"(HTTP {response.status}); upgrade that worker before running this feature"
-                        )
-                    capability = await response.json()
-                if not isinstance(capability, dict):
-                    raise RuntimeError(f"Image-list worker {worker_id} returned malformed capability data")
-                supported = capability.get("supported")
-                protocol_version = capability.get("protocol_version")
-                nodes = capability.get("node_ids")
-                if (
-                    not isinstance(supported, bool)
-                    or type(protocol_version) is not int
-                    or not isinstance(nodes, list)
-                    or any(not isinstance(node_id, str) for node_id in nodes)
-                ):
-                    raise RuntimeError(f"Image-list worker {worker_id} returned malformed capability data")
-                if (
-                    supported is not True
-                    or protocol_version != IMAGE_LIST_PROTOCOL_VERSION
-                    or not required_nodes.issubset(nodes)
-                ):
-                    raise RuntimeError(
-                        f"Image-list worker {worker_id} is incompatible: requires protocol "
-                        f"v{IMAGE_LIST_PROTOCOL_VERSION} and nodes "
-                        f"{sorted(required_nodes)}; no work was submitted"
-                    )
-
-                check = "queue"
-                queue_url = build_worker_url(worker, "/queue")
-                async with session.get(queue_url, timeout=aiohttp.ClientTimeout(total=timeout)) as response:
-                    if response.status != 200:
-                        raise RuntimeError(f"Image-list worker {worker_id} queue preflight returned HTTP {response.status}")
-                    queue = await response.json()
-                if (
-                    not isinstance(queue, dict)
-                    or not isinstance(queue.get("queue_running"), list)
-                    or not isinstance(queue.get("queue_pending"), list)
-                ):
-                    raise RuntimeError(f"Image-list worker {worker_id} returned malformed /queue state")
-                idle = not queue["queue_running"] and not queue["queue_pending"]
-                return worker_id, idle
-        except (asyncio.TimeoutError, aiohttp.ClientConnectionError) as exc:
-            log(
-                f"[Distributed] Image-list worker {worker_id} unavailable during {check} check; "
-                f"skipping: {exc}"
-            )
-            return worker_id, None
-
-    statuses = await asyncio.gather(*(inspect(worker) for worker in workers))
-    idle_workers = {
-        worker_id: True
-        for worker_id, idle in statuses
-        if idle is True
-    }
-    if require_idle and not idle_workers:
-        raise RuntimeError("No selected feature-compatible direct worker is idle; unrelated worker prompts were left untouched")
-    return idle_workers
 
 
 async def select_active_workers(
@@ -321,10 +203,6 @@ 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:

api/orchestration/media_sync.py

@@ -29,11 +29,6 @@ def _normalize_media_reference(value):
         return None
     cleaned = re.sub(r"\s*\[\w+\]$", "", value).strip().replace("\\", "/")
     if MEDIA_FILE_RE.search(cleaned):
-        # Output/temp annotations select a different master-side media root.
-        # Preserve them through discovery/loading and as distinct rewrite keys.
-        annotation = re.search(r"\s*\[(output|temp)\]$", value)
-        if annotation:
-            return f"{cleaned} [{annotation.group(1)}]"
         return cleaned
     return None
 
@@ -148,40 +143,10 @@ async def fetch_worker_path_separator(worker, trace_execution_id=None):
         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()
     normalized = filename.replace("\\", "/")
-    annotation = re.search(r"\s*\[(output|temp)\]$", normalized)
-    if annotation:
-        # Worker uploads live under input/. Do not collide across master roots,
-        # or overwrite a different reference with the same basename.
-        basename = normalized[:annotation.start()].rstrip().rsplit("/", 1)[-1]
-        normalized = f"distributed-media/{annotation.group(1)}/{file_hash}/{basename}"
 
     check_url = build_worker_url(worker, "/distributed/check_file")
     try:

api/orchestration/prompt_transform.py

@@ -328,19 +328,8 @@ def _find_delegate_master_retained_upstream_nodes(prompt_obj, start_ids):
     return connected
 
 
-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.
-    """
+def prune_prompt_for_worker(prompt_obj):
+    """Prune worker prompt to distributed nodes and their upstream dependencies."""
     collector_ids = find_nodes_by_class(prompt_obj, "DistributedCollector")
     upscale_ids = find_nodes_by_class(prompt_obj, "UltimateSDUpscaleDistributed")
     distributed_ids = collector_ids + upscale_ids
@@ -375,13 +364,7 @@ def prune_prompt_for_worker(prompt_obj, use_worker_sink=False):
                 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:
+            if image_connected:
                 preview_id = next_id()
                 pruned_prompt[preview_id] = {
                     "inputs": {"images": [dist_id, 0]},
@@ -458,55 +441,11 @@ def prepare_delegate_master_prompt(prompt_obj, collector_ids):
     return pruned_prompt
 
 
-def _resolve_direct_workers_for_image_list(config, requested_ids):
-    """Resolve enabled workers by default, or exact explicitly selected direct workers."""
-    requested = None if requested_ids is None else {str(worker_id) for worker_id in requested_ids}
-    configured_workers = (config or {}).get("workers", [])
-    configured_ids = {
-        str(worker.get("id") or "").strip()
-        for worker in configured_workers
-        if str(worker.get("id") or "").strip()
-    }
-    if requested is not None:
-        unknown_ids = requested - configured_ids
-        if unknown_ids:
-            raise RuntimeError(
-                "Unknown explicitly selected image-list worker ID(s): "
-                + ", ".join(sorted(unknown_ids))
-            )
-    direct = []
-    for worker in configured_workers:
-        worker_id = str(worker.get("id") or "").strip()
-        if not worker_id:
-            continue
-        if requested is None:
-            if not worker.get("enabled", False):
-                continue
-        elif worker_id not in requested:
-            continue
-        if str(worker.get("type") or "local").strip().lower() == "serverless":
-            continue
-        try:
-            port = int(worker.get("port", worker.get("listen_port", 8188)) or 8188)
-        except (TypeError, ValueError) as exc:
-            raise RuntimeError(f"Invalid port for image-list worker {worker_id}") from exc
-        direct.append({
-            "id": worker_id,
-            "name": str(worker.get("name") or worker_id),
-            "host": worker.get("host"),
-            "port": port,
-            "type": worker.get("type", "local"),
-        })
-    return direct
-
-
 def generate_job_id_map(prompt_index, prefix):
     """Create stable per-node job IDs for distributed nodes."""
     job_map = {}
-    distributed_nodes = (
-        prompt_index.nodes_for_class("DistributedCollector")
-        + prompt_index.nodes_for_class("UltimateSDUpscaleDistributed")
-        + prompt_index.nodes_for_class("DistributedImageListStart")
+    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}"
@@ -601,22 +540,6 @@ def _override_value_nodes(prompt_copy, prompt_index, is_master, participant_id,
             inputs["worker_id"] = f"worker_{worker_index_map.get(participant_id, 0)}"
 
 
-def _override_image_list_nodes(prompt_copy, prompt_index, participant_id, worker_ids,
-                               job_id_map, master_url, delegate_master, workflow_meta=None):
-    """Attach protocol context to Start without changing its visible V3 sockets."""
-    for node_id in prompt_index.nodes_for_class("DistributedImageListStart"):
-        node = prompt_copy.get(node_id)
-        if not isinstance(node, dict):
-            continue
-        inputs = node.setdefault("inputs", {})
-        inputs["image_list_enabled_worker_ids"] = json.dumps(list(worker_ids))
-        inputs["image_list_master_participates"] = participant_id == "master" and not delegate_master
-        inputs["image_list_master_callback"] = str(master_url or "")
-        inputs["image_list_execution_id"] = job_id_map.get(node_id, node_id)
-        if workflow_meta is not None:
-            inputs["image_list_workflow_meta"] = workflow_meta
-
-
 def apply_participant_overrides(
     prompt_copy,
     participant_id,
@@ -625,7 +548,6 @@ def apply_participant_overrides(
     master_url,
     delegate_master,
     prompt_index,
-    workflow_meta=None,
 ):
     """Return a prompt copy with hidden inputs configured for master/worker."""
     is_master = participant_id == "master"
@@ -653,15 +575,5 @@ def apply_participant_overrides(
         master_url,
         enabled_json,
     )
-    _override_image_list_nodes(
-        prompt_copy,
-        prompt_index,
-        participant_id,
-        enabled_worker_ids,
-        job_id_map,
-        master_url,
-        delegate_master,
-        workflow_meta=workflow_meta,
-    )
 
     return prompt_copy

api/orchestration/runpod_client.py

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

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

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

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

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

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

@@ -21,12 +21,7 @@ from .orchestration.dispatch import (
     select_active_workers,
     select_least_busy_worker,
 )
-from .orchestration.media_sync import (
-    convert_paths_for_platform,
-    fetch_worker_has_node,
-    fetch_worker_path_separator,
-    sync_worker_media,
-)
+from .orchestration.media_sync import convert_paths_for_platform, fetch_worker_path_separator, sync_worker_media
 from .orchestration.prompt_transform import (
     PromptIndex,
     apply_participant_overrides,
@@ -34,10 +29,7 @@ from .orchestration.prompt_transform import (
     generate_job_id_map,
     prepare_delegate_master_prompt,
     prune_prompt_for_worker,
-    WORKER_SINK_CLASS,
-    _resolve_direct_workers_for_image_list,
 )
-from ..image_list.runtime import validate_image_list_regions
 
 
 prompt_server = server.PromptServer.instance
@@ -82,14 +74,6 @@ 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)
@@ -170,21 +154,6 @@ 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,
@@ -193,21 +162,12 @@ async def _prepare_worker_payload(
 
         worker_type = str(worker.get("type") or "local").strip().lower()
         is_remote_like = bool(worker.get("host")) and worker_type != "local"
-        # 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()
+        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)
 
-        worker_prompt = prune_prompt_for_worker(worker_prompt, use_worker_sink=use_worker_sink)
+        worker_prompt = prune_prompt_for_worker(worker_prompt)
         worker_prompt = apply_participant_overrides(
             worker_prompt,
             worker["id"],
@@ -254,8 +214,6 @@ async def orchestrate_distributed_execution(
     execution_trace_id = trace_execution_id or _generate_execution_trace_id()
 
     config = load_config()
-    if delegate_master is None:
-        delegate_master = bool(config.get("settings", {}).get("master_delegate_only", False))
     use_websocket = bool(config.get("settings", {}).get("websocket_orchestration", False))
     master_url = build_master_url(config=config, prompt_server_instance=prompt_server)
     (
@@ -265,42 +223,8 @@ async def orchestrate_distributed_execution(
         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)
-    image_list_start_ids = prompt_index.nodes_for_class("DistributedImageListStart")
-    image_list_mode = bool(image_list_start_ids)
-    if image_list_mode:
-        # Validate before probing or queueing anything so malformed boundaries and
-        # old worker releases fail before any upstream image is materialized.
-        validate_image_list_regions(prompt_obj)
-        other_distributed = (
-            prompt_index.nodes_for_class("DistributedCollector")
-            + prompt_index.nodes_for_class("UltimateSDUpscaleDistributed")
-        )
-        if other_distributed:
-            raise ValueError(
-                "Distributed Image List v1 cannot be mixed with Distributed Collector or "
-                "Ultimate SD Upscale nodes in one prompt"
-            )
-        configured_serverless = [
-            str(worker.get("id")) for worker in config.get("workers", [])
-            if worker.get("enabled", False)
-            and str(worker.get("type") or "local").lower() == "serverless"
-            and (requested_ids is None or str(worker.get("id")) in {str(x) for x in requested_ids})
-        ]
-        if configured_serverless:
-            trace_debug(
-                execution_trace_id,
-                "Image List v1 explicitly excludes unsupported serverless pools: "
-                + ", ".join(configured_serverless),
-            )
-        workers = _resolve_direct_workers_for_image_list(config, requested_ids)
-        if not workers and delegate_master:
-            raise RuntimeError(
-                "Distributed Image List in orchestrator-only mode requires a selected direct worker; "
-                "serverless pools are unsupported"
-            )
-    else:
-        workers = _resolve_enabled_workers(config, requested_ids)
     load_balance_requested = _prompt_requests_load_balance(prompt_index)
     trace_debug(
         execution_trace_id,
@@ -316,37 +240,26 @@ async def orchestrate_distributed_execution(
         ),
     )
 
-    # Keep the legacy Collector fallback, but never override image-list
-    # orchestrator-only mode merely because its direct workers are unavailable.
-    if not image_list_mode and not workers and delegate_master:
+    # 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, selected_delegate_master = await select_active_workers(
+    active_workers, delegate_master = await select_active_workers(
         workers,
         use_websocket,
-        False if image_list_mode else delegate_master,
+        delegate_master,
         trace_execution_id=execution_trace_id,
         probe_concurrency=worker_probe_concurrency,
     )
-    if image_list_mode:
-        from .orchestration.dispatch import preflight_image_list_workers
-        if not active_workers and delegate_master:
-            raise RuntimeError(
-                "No selected direct worker is reachable for Distributed Image List in orchestrator-only mode"
-            )
-        idle_ids = (await preflight_image_list_workers(active_workers, require_idle=delegate_master)
-                    if active_workers else {})
-        active_workers = [worker for worker in active_workers if worker["id"] in idle_ids]
-        if not active_workers and delegate_master:
-            raise RuntimeError("No feature-compatible direct worker is idle in orchestrator-only mode")
-    else:
-        delegate_master = selected_delegate_master
 
-    if load_balance_requested and not image_list_mode:
+    if load_balance_requested:
         candidate_workers = list(active_workers)
         if not delegate_master:
             # Include master in load balancing only when master participation is enabled.
@@ -417,6 +330,9 @@ async def orchestrate_distributed_execution(
             queue_result.get("node_errors", {}),
         )
 
+    for job_id in job_id_map.values():
+        await _ensure_distributed_queue(job_id)
+
     master_prompt = prompt_index.copy_prompt()
     master_prompt = apply_participant_overrides(
         master_prompt,
@@ -426,10 +342,9 @@ async def orchestrate_distributed_execution(
         master_url,
         delegate_master,
         prompt_index,
-        workflow_meta=workflow_meta,
     )
 
-    if delegate_master and not image_list_mode:
+    if delegate_master:
         collector_ids = find_nodes_by_class(master_prompt, "DistributedCollector")
         upscale_nodes = find_nodes_by_class(master_prompt, "UltimateSDUpscaleDistributed")
         if upscale_nodes:
@@ -450,7 +365,7 @@ async def orchestrate_distributed_execution(
             + ", ".join(f"{worker['name']} ({worker['id']})" for worker in active_workers),
         )
     worker_payloads = []
-    if active_workers and not image_list_mode:
+    if active_workers:
         worker_prep_semaphore = asyncio.Semaphore(worker_prep_concurrency)
         media_sync_semaphore = asyncio.Semaphore(media_sync_concurrency)
         worker_payloads = await asyncio.gather(
@@ -472,20 +387,7 @@ async def orchestrate_distributed_execution(
             ]
         )
 
-    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.
-    if not image_list_mode:
-        for job_id in job_id_map.values():
-            await _ensure_distributed_queue(job_id)
-
-    if regular_payloads:
+    if worker_payloads:
         await asyncio.gather(
             *[
                 dispatch_worker_prompt(
@@ -496,42 +398,21 @@ async def orchestrate_distributed_execution(
                     use_websocket=use_websocket,
                     trace_execution_id=execution_trace_id,
                 )
-                for worker, wprompt in regular_payloads
+                for worker, wprompt in worker_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
+    queue_result = await queue_prompt_payload(
+        master_prompt,
+        workflow_meta,
+        client_id,
+        include_queue_metadata=True,
+    )
     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, prompt_number, (len(active_workers) if image_list_mode else len(worker_payloads)), node_errors
+    return prompt_id, prompt_number, len(worker_payloads), node_errors

api/serverless_routes.py

@@ -1,98 +0,0 @@
-"""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/worker_routes.py

@@ -16,14 +16,12 @@ 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 ..workers.ports import allocate_worker_ports
 from .schemas import require_fields, validate_worker_id
 from ..workers.detection import (
     get_machine_id,
@@ -279,13 +277,6 @@ def _get_cuda_info():
 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)
@@ -294,9 +285,7 @@ def _collect_network_info_sync():
         "all_ips": all_ips,
         "recommended_ip": recommended_ip,
         "cuda_device": cuda_device,
-        "cuda_device_count": device_count,
-        "master_port": master_port,
-        "local_worker_ports": worker_ports,
+        "cuda_device_count": physical_device_count if physical_device_count > 0 else cuda_device_count,
     }
 
 
@@ -460,9 +449,6 @@ async def launch_worker_endpoint(request):
         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
         
@@ -497,8 +483,6 @@ async def launch_worker_endpoint(request):
                 "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)
             
@@ -557,9 +541,6 @@ async def get_local_worker_status_endpoint(request):
         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"]:

conftest.py

@@ -1,7 +1,7 @@
 # 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
+# (from .distributed import ...) that fail when pytest tries to import it as a
 # standalone module during Package.setup() for the root package node.
 #
 # Fix: patch Package.setup() to skip the root-package's __init__.py import.

distributed.py

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

docs/compressed-result-transfer.md

@@ -1,38 +0,0 @@
-# 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/image-list-queue.md

@@ -1,99 +0,0 @@
-# Dynamic image-list queue
-
-Process a list of images across the master and direct local or remote ComfyUI workers without assigning fixed chunks to each participant.
-
-The master evaluates the input list once. By default, the master and each selected idle worker receive one image job at a time; when a participant finishes, it gets the next unassigned image. Faster participants can therefore process more images. The master restores input order before running downstream nodes.
-
-The existing **Orchestrator Only** setting (`settings.master_delegate_only`) excludes the master from per-image processing. The master still creates the input list and runs downstream nodes. API callers can override the setting for a run with `delegate_master`. There is no second image-list-specific toggle.
-
-## Wiring
-
-```text
-Image list or IMAGE batch
-          |
-Distributed Image List Start -------- context ---------+
-          | image                                     |
-          v                                           |
-    Per-image processing                              |
-          | IMAGE                                     |
-          v                                           v
-                 Distributed Image List End
-                            | image_list
-                            v
-                     Preview / Save / next stage
-```
-
-1. Install the same feature-capable version of Distributed on the master and every selected worker. Workers on an older image-list protocol fail preflight before any work is sent. Restart those instances and refresh the frontend after installation.
-2. Enable Distributed and select the **direct** workers to use. The master also participates unless **Orchestrator Only** is enabled. With the master participating, an empty worker selection runs locally through the same image-list scheduler; orchestrator-only mode requires at least one reachable, idle direct worker. The master must be reachable from remote workers at its configured callback address.
-3. Feed your image list or tensor batch into **Distributed Image List Start → images**.
-4. Connect its `image` output through the per-image processing nodes to **Distributed Image List End → images**.
-5. Connect Start's **context** directly to End's **context**. This is a dedicated `DISTRIBUTED_IMAGE_LIST_CONTEXT` socket, not editable text.
-6. Put preview/save nodes **after End**, then queue normally. API callers should submit through `/distributed/queue`, not directly through `/prompt`.
-
-ComfyUI represents these as ordinary `IMAGE` sockets with list semantics, not a separate `IMAGE_LIST` socket type. Outside the region there is a list; for each assigned item, Start provides one frame and its aligned values. End returns ordered lists, not a concatenated tensor batch, so differently sized images can be preserved.
-
-## Inputs and outputs
-
-Start accepts:
-
-- `images`: one image, an IMAGE batch, or a ComfyUI image list. A batch is split into individual frames; each frame becomes one job.
-- Optional `mask`.
-- Optional `value_1`, `value_2`, … inputs. Only `value_1` is shown at first; connecting it adds `value_2`, and so on, up to `value_10`. Each accepts an IMAGE, MASK, STRING, INT, FLOAT or BOOLEAN.
-
-`mask` and every value follow one rule:
-
-- **1 value** is shared by every image.
-- **One value per image** pairs them in order.
-- **Any other count** is an error. Values are never padded or repeated silently.
-
-Only `images` is split into frames. A batch connected to a value input stays a single value, so a 2-image batch can travel with each item (for example, two reference images per job). To pass a different batch to each image, connect a list with one batch per image.
-
-Start outputs, in order: `image`, `mask`, `item_index` (zero-based), `context`, then one `value_N` output per connected `value_N` input, with the same type. Unused value outputs are hidden. Seeds and strengths are ordinary values: connect a single number to share it, a list for one per image, or derive it from `item_index`.
-
-Values that cannot be transported exactly (models, latents, conditioning and other objects) are rejected. Load models inside the region from a loader connected directly to the processing nodes, as described below.
-
-End accepts `images`, optional `masks` / `texts`, and `context`. It returns `image_list`, `masks`, and `texts` in original item order. Each job must produce exactly one image, at most one mask, and at most one text value. Unconnected optional outputs contain `None` entries.
-
-Models are loaded and cached on each processing participant, not serialized between machines. Connect loader outputs directly to the processing nodes. The master needs the region's model files when participating, as does each selected worker. The master still validates the prompt: it needs the relevant node definitions and any model-name/file availability required by their validation (including combo lists). This is not a model-free-master guarantee.
-
-Per-image data must pass through Start; do not bypass it with upstream images, latents or generated batches connected directly into the processing region. Shared dependencies across Start are restricted to registered model outputs or scalar constants, not arbitrary image-producing leaf nodes.
-
-End's optional `texts` is a connection-only socket. For workflows saved before this socket became connection-only, recreate End and reconnect its wires if an old text widget or literal `texts: ""` remains in the compiled prompt. Arbitrary text literals are not accepted at End.
-
-File-based dependencies keep ComfyUI's `[output]` / `[temp]` root annotations and are read on the master. Their worker copies use separate content-addressed paths under `input/distributed-media/`; each worker's rewritten prompt is reused across items without reusing item tokens.
-
-## Scope and failure behavior
-
-- This version supports the master and direct workers, not Runpod serverless pools. Unsupported pools are excluded without contacting the provider for this job.
-- Worker eligibility is determined before execution and checked again before dispatch. Busy or transiently unreachable workers are excluded; healthy candidates can continue. Malformed capability/queue responses and incompatible versions fail explicitly. Workers joining later are not added to an already-running region.
-- Omitted worker IDs select enabled direct workers. Explicit IDs select the named configured direct workers, including disabled ones; an empty list stays empty, unknown IDs fail, and serverless workers are excluded. Local workers use the master's loopback callback; remote workers use its configured public address.
-- Only the enclosed processing is distributed. List creation and downstream saving run on the master in both modes. With participation enabled, the master is a normal scheduler participant and can run alone when no direct workers are available. **Orchestrator-only mode never silently falls back to master processing.** Failed items are not reassigned automatically.
-- Local items execute while the enclosing master prompt waits; they are not queued behind that prompt on the master's `/prompt` queue. The enclosed branch is not also evaluated as a full list by the outer prompt.
-- Start/End regions may be chained sequentially and share model-loader dependencies. Nested/overlapping per-image regions and data escaping a region without End are rejected. Collector and distributed-upscaler nodes cannot be mixed into the same prompt in this first version.
-- Output nodes inside a region are rejected because ComfyUI would execute them independently on the master. Avoid other side-effecting processing nodes as well; save after End.
-- The entire input list is materialized before scheduling; this is not streaming input. The limit is 4,096 frames. Remote transport preserves floating-point tensor dtype and values rather than encoding to lossy images. Each transported tensor is limited to 50 MiB, and the complete JSON/base64 HTTP request to 101 MiB; image/mask/value combinations can hit the aggregate request limit first. In-process master items use tensors directly and are not subject to those HTTP limits.
-- Results are accepted once per item using assignment tokens. Missing, duplicate, malformed and wrong-index results fail rather than silently returning a partial list. No failed or ambiguous submission is retried automatically.
-- Worker errors fail the region. The worker timeout measures loss of evidence for the owned prompt ID, not maximum render duration. Pending/running evidence refreshes liveness; unrelated busy jobs and failed polls do not. Healthy cold loads and long renders can continue.
-- User interruption is checked during submission and result waiting. Cleanup targets only this region's acknowledged prompt IDs and staged input tokens. Submission may wait briefly for an acknowledgement before cancellation; remote transport and cleanup have finite time limits. Local cleanup stops at node/progress boundaries and waits for its execution thread to finish before the parent prompt can resume. A synchronous node that does not cooperate with cancellation can delay this beyond the configured deadline. Lost acknowledgements are ambiguous, never retried automatically. Unreachable cleanup is logged as unconfirmed and may require checking the worker manually; cancellation does not guarantee remote termination through a broken network.
-- State is in memory. Restarting the master does not resume an interrupted region. Abandoned master contexts and expired worker staging are actively swept on the server loop; active/queued master work is protected from age-only expiry. End also removes its context on success, failure or interruption. This releases feature-owned references, not ComfyUI's separate execution cache.
-
-### Time and size limits
-
-`settings.image_list_job_timeout_seconds` defaults to `0` (no arbitrary total-run cap). A positive finite number opts into an overall deadline with owned-work cleanup; negative, non-finite or malformed settings fail explicitly. There is no fixed one-hour region deadline. The setting is separate from the existing worker liveness timeout.
-
-Uploads, media preparation and submission acknowledgements are independently bounded. WebSocket acknowledgement waiting includes connect/send/receive and is not extended by heartbeats or unrelated messages. A user interrupt allows up to one second for a late acknowledgement; cleanup requests use five-second bounds each.
-
-**4K float32 RGB outputs cannot be returned by direct workers under these transport limits.** A 3840×2160 RGB float32 tensor exceeds 50 MiB before base64 encoding. Mixed master/worker 1080p-to-4K upscaling therefore cannot use this configuration; master-only execution does not use the HTTP transport. Dtypes/values are never silently reduced to fit. Complete encoded input requests are checked before upload, including image/mask/value combinations; output limits are checked when the worker has computed the result. Reverse proxies may impose lower HTTP limits.
-
-Use trusted workers and trusted workflows, as with the existing ComfyUI API. Capability tokens protect assignment/result matching; they are not a replacement for network access control.
-
-## Verification
-
-The generic acceptance helper starts two isolated loopback worker test servers and uses actual CPU ComfyUI `PromptExecutor` instances. Production queue preparation, list nodes, graph extraction, HTTP item/result routes, dynamic scheduling, float transport, ordering and cancellation are exercised. The harness checks orchestrator-only, mixed master/worker and master-only execution, including exact-once item assignment, source-once materialization and cached local model-loader dependencies. The lightweight queue/history facade is test code; no managed ComfyUI service or model is started.
-
-```bash
-COMFYUI_SOURCE_ROOT=/path/to/ComfyUI /path/to/ComfyUI/.venv/bin/python -m pytest tests -q
-bash scripts/test-web.sh
-```
-
-Without `COMFYUI_SOURCE_ROOT`, real-framework tests are explicitly skipped. This coverage is not a substitute for checking your model workflow and connectivity on the actual GPU worker machines.

docs/serverless-workers.md

@@ -1,173 +0,0 @@
-# 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

@@ -1,53 +0,0 @@
-# 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,8 +74,6 @@ 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)

image_list/__init__.py

@@ -1 +0,0 @@
-"""Dynamic per-image image-list execution internals."""

image_list/lifecycle.py

@@ -1,142 +0,0 @@
-"""Server-owned background cleanup for image-list runtime state."""
-from __future__ import annotations
-
-import asyncio
-import logging
-import math
-import time
-from collections.abc import Iterable
-from typing import Any, cast
-
-LOGGER = logging.getLogger(__name__)
-DEFAULT_SWEEP_INTERVAL_SECONDS = 30.0
-DEFAULT_ORPHAN_GRACE_SECONDS = 300.0
-_LIFECYCLE_MARKER = object()
-
-
-def _queue_prompt_ids(prompt_server: Any) -> set[str]:
-    """Read running/pending prompt IDs, preferring ComfyUI's volatile view."""
-    queue = getattr(prompt_server, "prompt_queue", None)
-    readers = (
-        getattr(queue, "get_current_queue_volatile", None),
-        getattr(queue, "get_current_queue", None),
-    )
-    errors: list[Exception] = []
-    for reader in readers:
-        if not callable(reader):
-            continue
-        try:
-            snapshot = reader()
-            if not isinstance(snapshot, (tuple, list)) or len(snapshot) < 2:
-                raise TypeError("prompt queue snapshot must contain running and pending entries")
-            running, pending = snapshot[0], snapshot[1]
-            prompt_ids: set[str] = set()
-            for entries in (running, pending):
-                if not isinstance(entries, Iterable):
-                    raise TypeError("prompt queue entries must be iterable")
-                for entry in entries:
-                    if isinstance(entry, (tuple, list)) and len(entry) > 1:
-                        prompt_id = entry[1]
-                        if isinstance(prompt_id, str):
-                            prompt_ids.add(prompt_id)
-            return prompt_ids
-        except Exception as exc:
-            errors.append(exc)
-
-    if errors:
-        raise RuntimeError("unable to read the prompt queue") from errors[-1]
-    raise RuntimeError("prompt queue has no supported queue reader")
-
-
-def _sweep_orphan_contexts(prompt_server: Any, runtime: Any, orphan_grace_seconds: float) -> None:
-    queue_read_ok = True
-    try:
-        queued_prompt_ids = _queue_prompt_ids(prompt_server)
-    except Exception:
-        # A failed queue read must not make a submitted context look orphaned.
-        LOGGER.warning("Could not read prompt queue during image-list lifecycle sweep", exc_info=True)
-        queued_prompt_ids = set()
-        queue_read_ok = False
-
-    now = time.monotonic()
-    with runtime._LOCK:
-        for context_id, state in list(runtime._MASTER_JOBS.items()):
-            if getattr(state, "active", False):
-                continue
-            prompt_id = getattr(state, "prompt_id", None)
-            if prompt_id is not None:
-                if not queue_read_ok or prompt_id in queued_prompt_ids:
-                    continue
-            created_at = getattr(state, "created_at", None)
-            expired = (
-                isinstance(created_at, (int, float))
-                and not isinstance(created_at, bool)
-                and math.isfinite(created_at)
-                and now - created_at >= orphan_grace_seconds
-            )
-            if expired and runtime._MASTER_JOBS.get(context_id) is state:
-                runtime.remove_master_job(context_id)
-
-
-def _sweep_once(prompt_server: Any, runtime: Any, orphan_grace_seconds: float) -> None:
-    try:
-        runtime._sweep_worker_items()
-    except Exception:
-        LOGGER.exception("Could not sweep expired image-list worker staging")
-    _sweep_orphan_contexts(prompt_server, runtime, orphan_grace_seconds)
-
-
-async def _sweep_loop(
-    prompt_server: Any,
-    runtime: Any,
-    interval_seconds: float,
-    orphan_grace_seconds: float,
-) -> None:
-    while True:
-        try:
-            _sweep_once(prompt_server, runtime, orphan_grace_seconds)
-        except Exception:
-            # Keep the server-owned task alive if one pass encounters unexpected state.
-            LOGGER.exception("Image-list lifecycle sweep failed")
-        await asyncio.sleep(interval_seconds)
-
-
-def install_lifecycle(
-    prompt_server: Any,
-    runtime: Any,
-    *,
-    interval_seconds: float = DEFAULT_SWEEP_INTERVAL_SECONDS,
-    orphan_grace_seconds: float = DEFAULT_ORPHAN_GRACE_SECONDS,
-) -> None:
-    """Register one cancellable periodic sweep in the prompt server's app context."""
-    interval_seconds = float(interval_seconds)
-    orphan_grace_seconds = float(orphan_grace_seconds)
-    if not math.isfinite(interval_seconds) or interval_seconds <= 0:
-        raise ValueError("interval_seconds must be a positive finite value")
-    if not math.isfinite(orphan_grace_seconds) or orphan_grace_seconds < 0:
-        raise ValueError("orphan_grace_seconds must be a non-negative finite value")
-
-    app = getattr(prompt_server, "app", None)
-    cleanup_ctx = getattr(app, "cleanup_ctx", None)
-    if cleanup_ctx is None:
-        raise TypeError("prompt_server must expose an app with cleanup_ctx")
-    if any(getattr(factory, "_image_list_lifecycle_marker", None) is _LIFECYCLE_MARKER
-           for factory in cleanup_ctx):
-        return
-
-    async def lifecycle_context(_app):
-        task = asyncio.create_task(
-            _sweep_loop(prompt_server, runtime, interval_seconds, orphan_grace_seconds),
-            name="image-list-lifecycle-sweeper",
-        )
-        try:
-            yield
-        finally:
-            task.cancel()
-            try:
-                await task
-            except asyncio.CancelledError:
-                pass
-
-    cast(Any, lifecycle_context)._image_list_lifecycle_marker = _LIFECYCLE_MARKER
-    cleanup_ctx.append(lifecycle_context)

image_list/local_executor.py

@@ -1,184 +0,0 @@
-"""A queue-free Comfy graph slot, only while the synchronous outer End is blocked.
-
-Do NOT use PromptExecutor.execute[_async] here. Its prompt lifecycle clears the
-process interrupt, resets progress handlers, changes RAM-cache callbacks and
-starts/ends a model tracker. Those belong to the suspended parent executor.
-Instead use core's ExecutionList/execute (including lazy/async/expanded nodes)
-with private provider-free caches and a private server facade. Model outputs are
-registered with the existing parent tracker, whose lifetime includes this work.
-
-Core progress and GraphBuilder defaults are process-global. Borrow them only
-inside the owned execution thread, restore them in finally, and join that thread
-before End returns. Cancellation never clears or raises the global interrupt:
-it stops locally at node/progress boundaries. A non-cooperative synchronous node
-must finish before cleanup can return; abandoning its thread is not safe.
-"""
-from __future__ import annotations
-
-import asyncio
-from concurrent.futures import ThreadPoolExecutor
-import threading
-from types import SimpleNamespace
-
-from .runtime import ImageListCancelledError
-
-
-_EXECUTION_LOCK = threading.Lock()
-
-
-class LocalExecutionSlot:
-    def __init__(self, *, model_callback=None):
-        self._pool = ThreadPoolExecutor(max_workers=1, thread_name_prefix="image-list-master")
-        self._stop = threading.Event()
-        self._future = None
-        self._caches = None
-        self._model_callback = model_callback
-        self._closed = False
-        self._loop = None
-        self._task = None
-        self._state_lock = threading.Lock()
-
-    def submit(self, prompt, prompt_id, end_id, workflow_meta=None):
-        if self._closed or self._stop.is_set():
-            raise ImageListCancelledError("Local image-list slot has stopped")
-        if self._future is not None and not self._future.done():
-            raise RuntimeError("Local image-list slot already has an in-flight item")
-        self._future = self._pool.submit(self._run, prompt, prompt_id, end_id, workflow_meta)
-        return self._future
-
-    def request_stop(self):
-        self._stop.set()
-        # Wakes a graph waiting on an async node as well as stopping synchronous
-        # nodes at their next progress/between-node checkpoint.
-        with self._state_lock:
-            if self._loop is not None and self._task is not None:
-                self._loop.call_soon_threadsafe(self._task.cancel)
-
-    async def close(self):
-        self.request_stop()
-        try:
-            if self._future is not None:
-                # Do not cancel this Future: cancelling an asyncio wrapper cannot
-                # kill its running thread, and falsely makes cleanup look done.
-                waiter = asyncio.wrap_future(self._future)
-                while not waiter.done():
-                    try:
-                        await asyncio.shield(waiter)
-                    except asyncio.CancelledError:
-                        continue
-                    except Exception:
-                        break  # The scheduler propagates execution errors.
-                if waiter.done() and not waiter.cancelled():
-                    waiter.exception()
-        finally:
-            self._pool.shutdown(wait=True, cancel_futures=True)
-            self._caches = None
-            self._closed = True
-
-    def _run(self, prompt, prompt_id, end_id, workflow_meta):
-        # The normal Comfy prompt worker is single-threaded and blocked in End.
-        # Fail closed if another local slot is nevertheless executing, rather
-        # than interleave writes to core's process-global execution state.
-        if not _EXECUTION_LOCK.acquire(blocking=False):
-            raise RuntimeError("Another local image-list executor is already active")
-        try:
-            asyncio.run(self._execute(prompt, prompt_id, end_id, workflow_meta))
-        except asyncio.CancelledError as exc:
-            raise ImageListCancelledError("Local image-list execution cancelled") from exc
-        finally:
-            with self._state_lock:
-                self._loop = self._task = None
-            _EXECUTION_LOCK.release()
-
-    def _check_stopped(self):
-        import comfy.model_management
-        if self._stop.is_set() or comfy.model_management.processing_interrupted():
-            raise ImageListCancelledError("Local image-list execution interrupted")
-
-    async def _execute(self, prompt, prompt_id, end_id, workflow_meta):
-        import execution
-        import torch
-        import comfy.model_management
-        import comfy.utils
-        from comfy_execution import progress
-        from comfy_execution.graph_utils import GraphBuilder
-
-        self._check_stopped()
-        with self._state_lock:
-            self._loop, self._task = asyncio.get_running_loop(), asyncio.current_task()
-        dynamic_prompt = execution.DynamicPrompt(prompt)
-        private_progress = progress.ProgressRegistry(prompt_id, dynamic_prompt)
-        parent_progress = progress.global_progress_registry
-        parent_hook = comfy.utils.PROGRESS_BAR_HOOK
-        prefix_keys = ("_default_prefix_root", "_default_prefix_call_index", "_default_prefix_graph_index")
-        parent_prefix = {key: getattr(GraphBuilder, key) for key in prefix_keys}
-        progress.global_progress_registry = private_progress
-
-        def progress_hook(value, total, preview=None, prompt_id=None, node_id=None):
-            self._check_stopped()
-            # Never route a child progress/preview event through the parent's
-            # client or reset its handlers. Keep the private progress accurate.
-            if node_id is not None:
-                private_progress.update_progress(node_id, value, total, preview)
-
-        comfy.utils.PROGRESS_BAR_HOOK = progress_hook
-        try:
-            with torch.inference_mode():
-                if self._caches is None:
-                    outputs = execution.HierarchicalCache(execution.CacheKeySetInputSignature)
-                    objects = execution.HierarchicalCache(execution.CacheKeySetID)
-                    self._caches = SimpleNamespace(outputs=outputs, objects=objects, all=[outputs, objects])
-                caches = self._caches
-                changed = execution.IsChangedCache(prompt_id, dynamic_prompt, caches.outputs)
-                for cache in caches.all:
-                    await cache.set_prompt(dynamic_prompt, prompt.keys(), changed)
-                    cache.clean_unused()
-                execution_list = execution.ExecutionList(dynamic_prompt, caches.outputs, self._model_callback)
-                execution_list.add_node(end_id)
-                server = SimpleNamespace(client_id=None, last_node_id=None, send_sync=lambda *a, **kw: None)
-                assets = execution.NoAssets(execution.args)
-                extra = {"extra_pnginfo": {"workflow": workflow_meta}} if workflow_meta is not None else {}
-                executed, pending_subgraphs, pending_async, ui = set(), {}, {}, {}
-                while not execution_list.is_empty():
-                    self._check_stopped()
-                    node_id, error, exc = await execution_list.stage_node_execution()
-                    if error is not None:
-                        raise RuntimeError(f"Local image-list graph failed: {error}") from exc
-                    self._check_stopped()
-                    result, error, exc = await execution.execute(
-                        server, dynamic_prompt, caches, node_id, extra, executed, prompt_id,
-                        execution_list, pending_subgraphs, pending_async, ui, assets)
-                    if result == execution.ExecutionResult.FAILURE:
-                        if isinstance(exc, (comfy.model_management.InterruptProcessingException,
-                                            ImageListCancelledError)):
-                            raise ImageListCancelledError("Local image-list execution interrupted") from exc
-                        raise RuntimeError(f"Local image-list node {node_id} failed: {error}") from exc
-                    if result == execution.ExecutionResult.PENDING:
-                        execution_list.unstage_node_execution()
-                    else:
-                        execution_list.complete_node_execution()
-                self._check_stopped()
-        finally:
-            # Core execute creates async-node and unblock tasks. This is a private
-            # event loop, so every other task is owned here. Drain/cancel them
-            # before restoring globals, including when the graph failed early.
-            with self._state_lock:
-                self._task = None  # Stop requests must not interrupt this drain.
-            try:
-                while True:
-                    pending = [task for task in asyncio.all_tasks() if task is not asyncio.current_task()]
-                    if not pending:
-                        break
-                    for task in pending:
-                        task.cancel()
-                    drain = asyncio.gather(*pending, return_exceptions=True)
-                    while not drain.done():
-                        try:
-                            await asyncio.shield(drain)
-                        except asyncio.CancelledError:
-                            continue
-            finally:
-                progress.global_progress_registry = parent_progress
-                comfy.utils.PROGRESS_BAR_HOOK = parent_hook
-                for key, value in parent_prefix.items():
-                    setattr(GraphBuilder, key, value)

nodes/__init__.py

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

nodes/collector.py

@@ -12,21 +12,15 @@ from comfy.utils import ProgressBar
 
 from ..utils.logging import debug_log, log
 from ..utils.config import get_worker_timeout_seconds, load_config, is_master_delegate_only
-from ..utils.constants import HEARTBEAT_INTERVAL, TRANSFER_MAX_CHUNK_BYTES, TRANSFER_MAX_REQUEST_BYTES
+from ..utils.constants import HEARTBEAT_INTERVAL
 from ..utils.image import tensor_to_pil, pil_to_tensor, ensure_contiguous
 from ..utils.network import build_worker_url, get_client_session, probe_worker
 from ..utils.audio_payload import encode_audio_payload
 from ..utils.async_helpers import run_async_in_server_loop
-from ..utils import transfer_codec
 
 prompt_server = _server.PromptServer.instance
 
 
-def _encode_all_chunks(frames, fmt, quality, max_chunk_bytes, max_frame_bytes):
-    """Materialize every transfer chunk (runs in an executor)."""
-    return list(transfer_codec.encode_chunks(frames, fmt, quality, max_chunk_bytes, max_frame_bytes))
-
-
 class DistributedCollectorNode:
     INPUT_IS_LIST = True
     EMPTY_AUDIO = {"waveform": torch.zeros(1, 2, 1), "sample_rate": 44100}
@@ -46,23 +40,6 @@ class DistributedCollectorNode:
             "optional": {
                 "images": ("IMAGE",),
                 "audio": ("AUDIO",),
-                "transfer_format": (
-                    list(transfer_codec.TRANSFER_FORMATS),
-                    {
-                        "default": transfer_codec.DEFAULT_TRANSFER_FORMAT,
-                        "tooltip": transfer_codec.TRANSFER_FORMAT_TOOLTIP,
-                    },
-                ),
-                "transfer_quality": (
-                    "INT",
-                    {
-                        "default": transfer_codec.DEFAULT_TRANSFER_QUALITY,
-                        "min": 1,
-                        "max": 100,
-                        "step": 1,
-                        "tooltip": transfer_codec.TRANSFER_QUALITY_TOOLTIP,
-                    },
-                ),
             },
             "hidden": {
                 "multi_job_id": ("STRING", {"default": ""}),
@@ -88,6 +65,46 @@ class DistributedCollectorNode:
             return value[0]
         return value
 
+    @staticmethod
+    def _combine_image_tensors(tensors, sources):
+        """Validate BHWC batches, retaining alpha and deterministic input order."""
+        nonempty = []
+        reference_shape = None
+        reference_source = None
+        for tensor, source in zip(tensors, sources):
+            if not isinstance(tensor, torch.Tensor):
+                raise ValueError(f"Collector expected an IMAGE tensor from {source}")
+            shape = tuple(tensor.shape)
+            if tensor.ndim != 4 or shape[-1] not in (1, 3, 4) or min(shape[1:3]) <= 0:
+                raise ValueError(
+                    f"Collector received invalid IMAGE shape {shape} from {source}; "
+                    "expected [B, H, W, C] with 1, 3, or 4 channels"
+                )
+            # Empty master placeholders must not determine the worker image grid.
+            if shape[0] == 0:
+                continue
+            if reference_shape is None:
+                reference_shape, reference_source = shape, source
+            elif shape[1:3] != reference_shape[1:3]:
+                raise ValueError(
+                    "Collector cannot combine images with different spatial dimensions: "
+                    f"{reference_source} has shape {reference_shape}; {source} has shape {shape}"
+                )
+            nonempty.append(tensor)
+        if not nonempty:
+            raise ValueError("Collector cannot return an empty image batch")
+
+        channels = max(tensor.shape[-1] for tensor in nonempty)
+        normalized = []
+        for tensor in nonempty:
+            if tensor.shape[-1] == 1 and channels > 1:
+                tensor = tensor.expand(-1, -1, -1, 3)
+            if tensor.shape[-1] == 3 and channels == 4:
+                # RGB is opaque; never discard alpha from another participant.
+                tensor = torch.cat((tensor, torch.ones_like(tensor[..., :1])), dim=-1)
+            normalized.append(ensure_contiguous(tensor))
+        return ensure_contiguous(torch.cat(normalized, dim=0))
+
     def _normalize_images_input(self, images):
         """Collapse ComfyUI list IMAGE inputs into a normal batched IMAGE tensor."""
         if isinstance(images, (list, tuple)):
@@ -97,7 +114,9 @@ class DistributedCollectorNode:
                 raise TypeError("Collector expected IMAGE list items to be torch.Tensor instances")
             if len(images) == 1:
                 return ensure_contiguous(images[0])
-            return ensure_contiguous(torch.cat([ensure_contiguous(image) for image in images], dim=0))
+            return self._combine_image_tensors(
+                images, [f"IMAGE list item {index}" for index in range(len(images))]
+            )
         return ensure_contiguous(images)
 
     def _normalize_audio_input(self, audio):
@@ -127,7 +146,7 @@ class DistributedCollectorNode:
             return None
         return {"waveform": torch.cat(waveforms, dim=-1), "sample_rate": sample_rate}
 
-    def run(self, 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):
+    def run(self, 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):
         if images is not None:
             images = self._normalize_images_input(images)
         audio = self._normalize_audio_input(audio)
@@ -140,8 +159,6 @@ class DistributedCollectorNode:
         worker_id = self._unwrap_list_input(worker_id)
         pass_through = self._unwrap_list_input(pass_through)
         delegate_only = self._unwrap_list_input(delegate_only)
-        transfer_format = self._unwrap_list_input(transfer_format)
-        transfer_quality = self._unwrap_list_input(transfer_quality)
 
         remote_only_master = (
             bool(multi_job_id)
@@ -172,181 +189,12 @@ class DistributedCollectorNode:
                 worker_batch_size,
                 worker_id,
                 delegate_only,
-                transfer_format,
-                transfer_quality,
             )
         )
         return result
 
-    @staticmethod
-    def _get_serverless_manager():
-        """Return the serverless manager only for collector jobs with serverless slots."""
-        manager = getattr(prompt_server, "distributed_serverless", None)
-        if manager is not None:
-            return manager
-        from ..api.orchestration.serverless_runtime import get_manager
-        return get_manager()
-
-    async def send_batch_to_master(
-        self,
-        image_batch,
-        audio,
-        multi_job_id,
-        master_url,
-        worker_id,
-        transfer_format=transfer_codec.DEFAULT_TRANSFER_FORMAT,
-        transfer_quality=transfer_codec.DEFAULT_TRANSFER_QUALITY,
-    ):
+    async def send_batch_to_master(self, image_batch, audio, multi_job_id, master_url, worker_id):
         """Send an image batch, optionally with audio, or an audio-only completion."""
-        batch_size = 0 if image_batch is None else image_batch.shape[0]
-        fmt = self._resolve_transfer_format(transfer_format) if batch_size else transfer_codec.LEGACY_FORMAT
-        if fmt != transfer_codec.LEGACY_FORMAT:
-            sent = await self._send_compressed_batch(
-                image_batch, audio, multi_job_id, master_url, worker_id, fmt, transfer_quality
-            )
-            if sent:
-                return
-            log(
-                "Worker - Master did not accept compressed transfer "
-                "(older ComfyUI-Distributed on the master?). Falling back to legacy_png."
-            )
-        await self._send_legacy_batch(image_batch, audio, multi_job_id, master_url, worker_id)
-
-    @staticmethod
-    def _resolve_transfer_format(transfer_format):
-        try:
-            fmt = transfer_codec.normalize_format(transfer_format)
-        except transfer_codec.TransferCodecError as exc:
-            log(f"Worker - {exc}; using legacy_png.")
-            return transfer_codec.LEGACY_FORMAT
-        if fmt not in transfer_codec.VIDEO_FORMATS or transfer_codec.video_codec_available(fmt):
-            return fmt
-        fallback = "h264" if fmt != "ffv1" and transfer_codec.video_codec_available("h264") else "png"
-        log(f"Worker - Encoder for '{fmt}' is not available in this PyAV build; using '{fallback}'.")
-        return fallback
-
-    @staticmethod
-    def _build_multipart(envelope, parts):
-        form = aiohttp.FormData()
-        form.add_field("envelope", json.dumps(envelope), content_type="application/json")
-        for i, data in enumerate(parts):
-            form.add_field("media", data, filename=f"{i}.bin", content_type="application/octet-stream")
-        return form
-
-    # Room left in each request for the multipart envelope (ids, frame
-    # metadata, part headers). Audio is budgeted separately.
-    TRANSFER_ENVELOPE_RESERVE_BYTES = 256 * 1024
-    # Body of a 400 from a master without compressed-transfer support: it
-    # tries request.json() on the multipart body and fails.
-    _OLD_MASTER_REJECTION = "Invalid JSON payload"
-
-    async def _send_compressed_batch(self, image_batch, audio, multi_job_id, master_url, worker_id, fmt, quality):
-        """Send the batch as compressed multipart chunks.
-
-        All chunks are encoded before anything is sent, so a frame that cannot
-        fit in one request fails the transfer without delivering a partial
-        batch. Every request stays within TRANSFER_MAX_REQUEST_BYTES.
-
-        Returns False (without having delivered anything) when the master
-        rejects the first chunk because it predates compressed transfer, so
-        the caller can retry with legacy_png. Any other rejection raises.
-        """
-        encoded_audio = encode_audio_payload(audio)
-        session = await get_client_session()
-        url = f"{master_url}/distributed/job_complete"
-        loop = asyncio.get_running_loop()
-
-        frames = await loop.run_in_executor(None, transfer_codec.tensor_batch_to_uint8, image_batch)
-        raw_bytes = int(frames.nbytes)
-        total = int(frames.shape[0])
-        media_limit = max(1, TRANSFER_MAX_REQUEST_BYTES - self.TRANSFER_ENVELOPE_RESERVE_BYTES)
-        started = time.time()
-        # Encoding is CPU-bound (seconds for video); keep the server loop responsive.
-        chunks = await loop.run_in_executor(
-            None, _encode_all_chunks, frames, fmt, quality, min(TRANSFER_MAX_CHUNK_BYTES, media_limit), media_limit,
-        )
-        del frames
-
-        # Attach audio to the final chunk only if the request stays in bounds;
-        # otherwise finish with a separate audio-only completion.
-        audio_bytes = len(json.dumps(encoded_audio)) if encoded_audio is not None else 0
-        last_media = sum(len(p) for p in chunks[-1]["parts"])
-        audio_inline = encoded_audio is not None and last_media + audio_bytes <= media_limit
-        audio_separate = encoded_audio is not None and not audio_inline
-
-        sent_bytes = 0
-        for index, chunk in enumerate(chunks):
-            parts = chunk.pop("parts")
-            is_final_chunk = index == len(chunks) - 1
-            envelope = {
-                "job_id": str(multi_job_id),
-                "worker_id": str(worker_id),
-                "is_last": bool(is_final_chunk and not audio_separate),
-                **chunk,
-            }
-            if is_final_chunk and audio_inline:
-                envelope["audio"] = encoded_audio
-            chunk_bytes = sum(len(p) for p in parts)
-            timeout_seconds = min(600, 60 + chunk_bytes // (256 * 1024))
-            try:
-                async with session.post(
-                    url,
-                    data=self._build_multipart(envelope, parts),
-                    timeout=aiohttp.ClientTimeout(total=timeout_seconds),
-                ) as response:
-                    if response.status >= 400:
-                        body = await response.text()
-                        if index == 0 and response.status == 400 and self._OLD_MASTER_REJECTION in body:
-                            debug_log(f"Worker - Master rejected compressed chunk: {body[:300]}")
-                            return False
-                        raise RuntimeError(
-                            f"Master rejected {fmt} chunk {index + 1}/{len(chunks)} "
-                            f"({chunk_bytes} bytes) with HTTP {response.status}: {body[:300]}. "
-                            "If this is a size limit, make COMFYUI_DISTRIBUTED_TRANSFER_MAX_REQUEST_BYTES "
-                            "on the worker smaller than the master's --max-upload-size."
-                        )
-            except Exception as e:
-                log(f"Worker - Failed to send {fmt} chunk to master: {e}")
-                debug_log(f"Worker - Full error details: URL={url}")
-                raise
-            sent_bytes += chunk_bytes
-
-        requests = len(chunks)
-        if audio_separate:
-            await self._post_json_completion(session, url, {
-                "job_id": str(multi_job_id),
-                "worker_id": str(worker_id),
-                "batch_idx": 0,
-                "audio": encoded_audio,
-                "is_last": True,
-            })
-            requests += 1
-
-        log(
-            f"Worker - Sent {total} image(s) as {fmt} in {requests} request(s): "
-            f"{sent_bytes / 1e6:.2f} MB (raw {raw_bytes / 1e6:.2f} MB, "
-            f"{transfer_codec.describe_ratio(sent_bytes, raw_bytes)} smaller) in {time.time() - started:.1f}s"
-        )
-        return True
-
-    @staticmethod
-    async def _post_json_completion(session, url, payload):
-        timeout_seconds = 60 if "image" in payload else 600
-        try:
-            async with session.post(
-                url,
-                json=payload,
-                timeout=aiohttp.ClientTimeout(total=timeout_seconds),
-            ) as response:
-                response.raise_for_status()
-        except Exception as e:
-            media_type = "image/audio" if "image" in payload else "audio-only"
-            log(f"Worker - Failed to send canonical {media_type} envelope to master: {e}")
-            debug_log(f"Worker - Full error details: URL={url}")
-            raise
-
-    async def _send_legacy_batch(self, image_batch, audio, multi_job_id, master_url, worker_id):
-        """Original transfer: one JSON request per frame with a base64 PNG."""
         encoded_audio = encode_audio_payload(audio)
         session = await get_client_session()
         url = f"{master_url}/distributed/job_complete"
@@ -383,7 +231,19 @@ class DistributedCollectorNode:
                 payloads.append(payload)
 
         for payload in payloads:
-            await self._post_json_completion(session, url, payload)
+            timeout_seconds = 60 if "image" in payload else 600
+            try:
+                async with session.post(
+                    url,
+                    json=payload,
+                    timeout=aiohttp.ClientTimeout(total=timeout_seconds),
+                ) as response:
+                    response.raise_for_status()
+            except Exception as e:
+                media_type = "image/audio" if "image" in payload else "audio-only"
+                log(f"Worker - Failed to send canonical {media_type} envelope to master: {e}")
+                debug_log(f"Worker - Full error details: URL={url}")
+                raise
 
     def _combine_audio(self, master_audio, worker_audio, empty_audio, worker_order=None):
         """Combine audio from master and workers into a single audio output.
@@ -468,9 +328,11 @@ class DistributedCollectorNode:
     ):
         """Assemble final tensor, or return None when the job contains only audio."""
         ordered_tensors = []
+        sources = []
         if not delegate_mode and images_on_cpu is not None:
             for i in range(master_batch_size):
                 ordered_tensors.append(images_on_cpu[i:i+1])
+                sources.append(f"master image {i}")
 
         ordered_worker_ids = [str(worker_id) for worker_id in (worker_order or [])]
         seen = set()
@@ -480,6 +342,7 @@ class DistributedCollectorNode:
                 continue
             for idx in sorted(worker_images[worker_id_str].keys()):
                 ordered_tensors.append(worker_images[worker_id_str][idx])
+                sources.append(f"worker {worker_id_str!r} image {idx}")
 
         # Append any unexpected worker ids deterministically.
         for worker_id_str in sorted(worker_images.keys()):
@@ -487,6 +350,7 @@ class DistributedCollectorNode:
                 continue
             for idx in sorted(worker_images[worker_id_str].keys()):
                 ordered_tensors.append(worker_images[worker_id_str][idx])
+                sources.append(f"worker {worker_id_str!r} image {idx}")
 
         cpu_tensors = []
         for t in ordered_tensors:
@@ -496,19 +360,17 @@ class DistributedCollectorNode:
             cpu_tensors.append(t)
 
         if cpu_tensors:
-            return ensure_contiguous(torch.cat(cpu_tensors, dim=0))
+            return self._combine_image_tensors(cpu_tensors, sources)
         if fallback_images is not None:
-            return ensure_contiguous(fallback_images)
+            return self._combine_image_tensors([fallback_images], ["master fallback"])
         return None
 
-    async def execute(self, images, audio, load_balance=False, multi_job_id="", is_worker=False, master_url="", enabled_worker_ids="[]", worker_batch_size=1, worker_id="", delegate_only=False, transfer_format=transfer_codec.DEFAULT_TRANSFER_FORMAT, transfer_quality=transfer_codec.DEFAULT_TRANSFER_QUALITY):
+    async def execute(self, images, audio, load_balance=False, multi_job_id="", is_worker=False, master_url="", enabled_worker_ids="[]", worker_batch_size=1, worker_id="", delegate_only=False):
         if is_worker:
             # Worker mode: send images and audio to master in a single batch
             image_count = 0 if images is None else images.shape[0]
             debug_log(f"Worker - Job {multi_job_id} complete. Sending {image_count} image(s) to master")
-            await self.send_batch_to_master(
-                images, audio, multi_job_id, master_url, worker_id, transfer_format, transfer_quality
-            )
+            await self.send_batch_to_master(images, audio, multi_job_id, master_url, worker_id)
             return (images, audio if audio is not None else self.EMPTY_AUDIO)
         else:
             delegate_mode = delegate_only or is_master_delegate_only()
@@ -523,17 +385,6 @@ class DistributedCollectorNode:
                 seen_enabled.add(worker_id_str)
                 enabled_workers.append(worker_id_str)
             expected_workers = set(enabled_workers)
-            serverless_workers = {
-                worker_id for worker_id in expected_workers if worker_id.startswith("serverless:")
-            }
-            serverless_manager = None
-
-            def get_serverless_manager():
-                nonlocal serverless_manager
-                if serverless_workers and serverless_manager is None:
-                    serverless_manager = self._get_serverless_manager()
-                return serverless_manager
-
             num_workers = len(expected_workers)
             if num_workers == 0:
                 return (images, audio if audio is not None else self.EMPTY_AUDIO)
@@ -600,177 +451,153 @@ class DistributedCollectorNode:
                 workers_done.add(done_worker_id)
                 p.update(1)  # +1 per completed expected worker
 
-            def consume_result(result):
-                nonlocal collected_count
-                if "serverless_error" in result:
-                    error = result.get("serverless_error")
-                    message = str(error).strip() if error is not None else "unspecified serverless failure"
-                    participant = result.get("worker_id", "unknown")
-                    raise RuntimeError(f"Serverless participant {participant} failed: {message}")
-
-                result_worker_id = str(result["worker_id"])
-                is_last = result.get("is_last", False)
-                collected_count += self._store_worker_result(worker_images, result)
-                debug_log(
-                    f"Master - Got canonical result from worker {result_worker_id}, "
-                    f"image {result.get('image_index', 0)}, is_last={is_last}"
-                )
-                result_audio = result.get("audio")
-                if result_audio is not None:
-                    worker_audio[result_worker_id] = result_audio
-                    debug_log(f"Master - Got audio from worker {result_worker_id}")
-                if is_last:
-                    mark_worker_done(result_worker_id)
-
             try:
                 while len(workers_done) < num_workers:
-                    # Check for user interruption to abort collection promptly.
+                    # Check for user interruption to abort collection promptly
                     comfy.model_management.throw_exception_if_processing_interrupted()
                     try:
+                        # Get the queue again each time to ensure we have the right reference
                         async with prompt_server.distributed_jobs_lock:
                             q = prompt_server.distributed_pending_jobs[multi_job_id]
+                            current_size = q.qsize()
+                        
                         result = await asyncio.wait_for(q.get(), timeout=slice_timeout)
-                        consume_result(result)
+                        worker_id = result['worker_id']
+                        is_last = result.get('is_last', False)
+                        count = self._store_worker_result(worker_images, result)
+                        collected_count += count
+                        debug_log(
+                            f"Master - Got canonical result from worker {worker_id}, "
+                            f"image {result.get('image_index', 0)}, is_last={is_last}"
+                        )
 
-                        # Real results refresh the regular-worker timeout baseline.
+                        # Collect audio data if present
+                        result_audio = result.get('audio')
+                        if result_audio is not None:
+                            worker_audio[worker_id] = result_audio
+                            debug_log(f"Master - Got audio from worker {worker_id}")
+
+                        # Record activity and refresh timeout baseline
                         last_activity = time.time()
                         base_timeout = float(get_worker_timeout_seconds())
 
+                        if is_last:
+                            mark_worker_done(worker_id)
+                        
                     except asyncio.TimeoutError:
+                        # If we still have time, continue polling; otherwise handle timeout
                         if (time.time() - last_activity) < base_timeout:
                             comfy.model_management.throw_exception_if_processing_interrupted()
                             continue
+                        # Re-check for user interruption after timeout expiry
                         comfy.model_management.throw_exception_if_processing_interrupted()
-                        missing_workers = expected_workers - workers_done
-
-                        # A queue delivery can race the timed wait. Process it before
-                        # declaring an owned provider slot stopped without a result.
-                        async with prompt_server.distributed_jobs_lock:
-                            current_q = prompt_server.distributed_pending_jobs.get(multi_job_id)
-                            queue_has_results = current_q is not None and not current_q.empty()
-                        if queue_has_results:
-                            continue
-
-                        missing_serverless = missing_workers & serverless_workers
-                        if missing_serverless:
-                            manager = get_serverless_manager()
-                            if manager is None:
-                                raise RuntimeError("Serverless manager is unavailable for pending collector slots")
-                            pending_serverless = []
-                            for participant_id in sorted(missing_serverless):
-                                if manager.is_pending(multi_job_id, participant_id):
-                                    pending_serverless.append(participant_id)
-                                elif manager.owns(multi_job_id, participant_id):
-                                    raise RuntimeError(
-                                        f"Serverless participant {participant_id} stopped without delivering a result"
-                                    )
-                            # Provider deadlines are managed by the serverless runtime;
-                            # do not apply the regular worker heartbeat timeout to it.
-                            if pending_serverless:
-                                continue
-
-                        missing_regular = missing_workers - serverless_workers
+                        missing_workers = set(str(w) for w in enabled_workers) - workers_done
                         elapsed = time.time() - last_activity
-                        for missing_worker_id in sorted(missing_regular):
+                        for missing_worker_id in sorted(missing_workers):
                             log(
                                 "Master - Heartbeat timeout: "
                                 f"worker={missing_worker_id}, elapsed={elapsed:.1f}s"
                             )
-                        if missing_regular:
-                            log(
-                                f"Master - Heartbeat timeout. Still waiting for workers: {list(sorted(missing_regular))} "
-                                f"(elapsed={elapsed:.1f}s)"
-                            )
+                        log(
+                            f"Master - Heartbeat timeout. Still waiting for workers: {list(missing_workers)} "
+                            f"(elapsed={elapsed:.1f}s)"
+                        )
 
-                        # Probe only regular worker /prompt endpoints.
+                        # Probe missing workers' /prompt endpoints to check if they are actively processing
                         any_busy = False
-                        if missing_regular:
-                            try:
-                                cfg = load_config()
-                                cfg_workers = cfg.get('workers', [])
-                                for wid in sorted(missing_regular):
-                                    wrec = next((w for w in cfg_workers if str(w.get('id')) == str(wid)), None)
-                                    if not wrec:
-                                        debug_log(f"Collector probe: worker {wid} not found in config")
-                                        continue
-                                    worker_url = build_worker_url(wrec)
-                                    try:
-                                        payload = await probe_worker(worker_url, timeout=2.0)
-                                        queue_remaining = None
-                                        if payload is not None:
-                                            queue_remaining = int(payload.get('exec_info', {}).get('queue_remaining', 0))
-                                        debug_log(
-                                            "Collector probe: worker "
-                                            f"{wid} online={payload is not None} queue_remaining={queue_remaining}"
+                        try:
+                            cfg = load_config()
+                            cfg_workers = cfg.get('workers', [])
+                            for wid in list(missing_workers):
+                                wrec = next((w for w in cfg_workers if str(w.get('id')) == str(wid)), None)
+                                if not wrec:
+                                    debug_log(f"Collector probe: worker {wid} not found in config")
+                                    continue
+                                worker_url = build_worker_url(wrec)
+                                try:
+                                    payload = await probe_worker(worker_url, timeout=2.0)
+                                    queue_remaining = None
+                                    if payload is not None:
+                                        queue_remaining = int(payload.get('exec_info', {}).get('queue_remaining', 0))
+                                    debug_log(
+                                        "Collector probe: worker "
+                                        f"{wid} online={payload is not None} queue_remaining={queue_remaining}"
+                                    )
+                                    if payload is not None and queue_remaining and queue_remaining > 0:
+                                        any_busy = True
+                                        log(
+                                            f"Master - Probe grace: worker {wid} appears busy "
+                                            f"(queue_remaining={queue_remaining}). Continuing to wait."
                                         )
-                                        if payload is not None and queue_remaining and queue_remaining > 0:
-                                            any_busy = True
-                                            log(
-                                                f"Master - Probe grace: worker {wid} appears busy "
-                                                f"(queue_remaining={queue_remaining}). Continuing to wait."
-                                            )
-                                            break
-                                    except Exception as e:
-                                        debug_log(f"Collector probe failed for worker {wid}: {e}")
-                            except Exception as e:
-                                debug_log(f"Collector probe setup error: {e}")
+                                        break
+                                except Exception as e:
+                                    debug_log(f"Collector probe failed for worker {wid}: {e}")
+                        except Exception as e:
+                            debug_log(f"Collector probe setup error: {e}")
 
                         if any_busy:
+                            # Refresh last_activity and continue waiting
                             last_activity = time.time()
+                            # Refresh base timeout in case the user changed it in UI
                             base_timeout = float(get_worker_timeout_seconds())
                             continue
-
+                        
+                        # Check queue size again with lock
                         async with prompt_server.distributed_jobs_lock:
-                            final_q = prompt_server.distributed_pending_jobs.get(multi_job_id)
-                            remaining_items = []
-                            if final_q is not None:
+                            if multi_job_id in prompt_server.distributed_pending_jobs:
+                                final_q = prompt_server.distributed_pending_jobs[multi_job_id]
+                                final_size = final_q.qsize()
+                                
+                                # Try to drain any remaining items
+                                remaining_items = []
                                 while not final_q.empty():
                                     try:
-                                        remaining_items.append(final_q.get_nowait())
+                                        item = final_q.get_nowait()
+                                        remaining_items.append(item)
                                     except asyncio.QueueEmpty:
                                         break
-                        for item in remaining_items:
-                            consume_result(item)
+                                
+                                if remaining_items:
+                                    # Process them
+                                    for item in remaining_items:
+                                        worker_id = item['worker_id']
+                                        is_last = item.get('is_last', False)
+
+                                        collected_count += self._store_worker_result(worker_images, item)
+                                        
+                                        if is_last:
+                                            mark_worker_done(worker_id)
+                            else:
+                                log(f"Master - Queue {multi_job_id} no longer exists!")
                         break
-
-                total_collected = sum(len(imgs) for imgs in worker_images.values())
-
-                # Remove the queue before combining so a late provider callback cannot
-                # resurrect results for a collector that has completed.
+            except comfy.model_management.InterruptProcessingException:
+                # Cleanup queue on interruption and re-raise to abort prompt cleanly
                 async with prompt_server.distributed_jobs_lock:
-                    prompt_server.distributed_pending_jobs.pop(multi_job_id, None)
+                    if multi_job_id in prompt_server.distributed_pending_jobs:
+                        del prompt_server.distributed_pending_jobs[multi_job_id]
+                raise
+            
+            total_collected = sum(len(imgs) for imgs in worker_images.values())
+            
+            # Clean up job queue
+            async with prompt_server.distributed_jobs_lock:
+                if multi_job_id in prompt_server.distributed_pending_jobs:
+                    del prompt_server.distributed_pending_jobs[multi_job_id]
 
-                combined_audio = self._combine_audio(master_audio, worker_audio, self.EMPTY_AUDIO, enabled_workers)
-                serverless_results = set(worker_images) & serverless_workers
-                try:
-                    combined = self._reorder_and_combine_tensors(
-                        worker_images, enabled_workers, master_batch_size, images_on_cpu, delegate_mode, images
-                    )
-                    if combined is None:
-                        debug_log(f"Master - Job {multi_job_id} complete with audio only")
-                    else:
-                        debug_log(f"Master - Job {multi_job_id} complete. Combined {combined.shape[0]} images total "
-                                  f"(master: {master_batch_size}, workers: {combined.shape[0] - master_batch_size})")
-                    return (combined, combined_audio)
-                except Exception as e:
-                    if serverless_results:
-                        raise RuntimeError(
-                            "Serverless result batch shape mismatch while combining participants "
-                            f"{sorted(serverless_results)}: {e}"
-                        ) from e
-                    log(f"Master - Error combining images: {e}")
-                    # Preserve collected audio for regular workers when image assembly fails.
-                    return (images, combined_audio)
-            except BaseException:
-                # Remove the queue first, preventing cancelled work from delivering into
-                # a collector that is already unwinding.
-                async with prompt_server.distributed_jobs_lock:
-                    prompt_server.distributed_pending_jobs.pop(multi_job_id, None)
-                if serverless_workers:
-                    try:
-                        manager = get_serverless_manager()
-                        if manager is not None:
-                            await manager.cancel(collector_job_id=multi_job_id)
-                    except Exception as cancel_error:
-                        debug_log(f"Collector - Serverless cancellation failed: {cancel_error}")
+            combined_audio = self._combine_audio(master_audio, worker_audio, self.EMPTY_AUDIO, enabled_workers)
+            try:
+                combined = self._reorder_and_combine_tensors(
+                    worker_images, enabled_workers, master_batch_size, images_on_cpu, delegate_mode, images
+                )
+                if combined is None:
+                    debug_log(f"Master - Job {multi_job_id} complete with audio only")
+                else:
+                    debug_log(f"Master - Job {multi_job_id} complete. Combined {combined.shape[0]} images total "
+                              f"(master: {master_batch_size}, workers: {combined.shape[0] - master_batch_size})")
+
+                return (combined, combined_audio)
+            except Exception as e:
+                log(f"Master - Error combining images: {e}")
+                # A fallback can silently lose worker results or feed an empty
+                # batch into SaveImage. Surface the original collector error.
                 raise

nodes/distributed_upscale.py

@@ -268,3 +268,12 @@ class UltimateSDUpscaleDistributed(
 
 # Ensure initialization before registering routes
 ensure_tile_jobs_initialized()
+
+# Node registration
+NODE_CLASS_MAPPINGS = {
+    "UltimateSDUpscaleDistributed": UltimateSDUpscaleDistributed,
+}
+
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "UltimateSDUpscaleDistributed": "Ultimate SD Upscale Distributed (No Upscale)",
+}

nodes/image_list.py

@@ -1,252 +0,0 @@
-"""Public V3 Start/End nodes for dynamic per-image image-list execution."""
-from __future__ import annotations
-
-import json
-
-from comfy_api.v0_0_2 import io
-
-from ..utils.async_helpers import run_async_in_server_loop
-from ..image_list.runtime import (
-    MAX_VALUE_SLOTS,
-    VALUE_NAMES,
-    VALUES_INPUT,
-    create_master_job,
-    normalize_image_list_inputs,
-    normalize_worker_result,
-    run_image_list_job,
-    encode_result,
-    encode_request_payload,
-    unwrap_single,
-)
-
-
-# Types the transport can carry exactly. Anything else (models, latents, ...)
-# must be loaded inside the region or shared as a model dependency instead.
-_VALUE_TYPES = [io.Image, io.Mask, io.String, io.Int, io.Float, io.Boolean]
-assert len(VALUE_NAMES) == MAX_VALUE_SLOTS
-
-
-class DistributedImageListStart(io.ComfyNode):
-    """Materialize a batch once, then expose ordered per-item typed values."""
-
-    @classmethod
-    def define_schema(cls):
-        return io.Schema(
-            node_id="DistributedImageListStart",
-            display_name="Distributed Image List Start",
-            category="image/distributed",
-            description=("Starts one per-image processing region. Each image becomes one job. "
-                         "mask and value_N inputs take 1 value (shared by every image) or exactly "
-                         "one value per image, paired in order."),
-            inputs=[
-                io.Image.Input("images", optional=True, tooltip="Image list or image batch; each frame becomes one job."),
-                io.Mask.Input("mask", optional=True,
-                              tooltip="1 mask shared by every image, or one mask per image."),
-                io.Autogrow.Input(VALUES_INPUT, optional=True, template=io.Autogrow.TemplateNames(
-                    io.MultiType.Input("value", types=_VALUE_TYPES,
-                                       tooltip=("Any image, mask, text, number or boolean. 1 value is shared by "
-                                                "every image; one value per image pairs them in order. "
-                                                "A batch stays one value.")),
-                    names=list(VALUE_NAMES), min=0)),
-            ],
-            outputs=[
-                io.Image.Output(display_name="image", is_output_list=True),
-                io.Mask.Output(display_name="mask", is_output_list=True),
-                io.Int.Output(display_name="item_index", is_output_list=True),
-                io.Custom("DISTRIBUTED_IMAGE_LIST_CONTEXT").Output(display_name="context", is_output_list=True),
-                # value_N outputs mirror value_N inputs. The web extension hides
-                # unused ones and shows the connected input's type.
-                *[io.AnyType.Output(id=name, display_name=name, is_output_list=True) for name in VALUE_NAMES],
-            ],
-            is_input_list=True,
-            accept_all_inputs=True,
-            not_idempotent=True,
-            hidden=[io.Hidden.prompt, io.Hidden.unique_id],
-        )
-
-    @classmethod
-    def execute(cls, images=None, mask=None, values=None,
-                worker_item_token="", worker_context_id="", image_list_enabled_worker_ids="[]",
-                image_list_master_callback="", image_list_execution_id="",
-                image_list_workflow_meta=None, image_list_master_participates=None,
-                image_list_local_execution=False):
-        worker_item_token = unwrap_single(worker_item_token)
-        worker_context_id = unwrap_single(worker_context_id)
-        image_list_workflow_meta = unwrap_single(image_list_workflow_meta)
-        if worker_item_token:
-            if unwrap_single(image_list_local_execution):
-                from ..image_list.runtime import consume_worker_item
-                item = consume_worker_item(worker_item_token)
-            else:
-                item = run_async_in_server_loop(_consume_worker_item(worker_item_token), timeout=10)
-            context_id = worker_context_id
-            if not isinstance(context_id, str) or not context_id:
-                raise RuntimeError("Worker image-list prompt is missing its paired context")
-            item_values = item.get("values") or {}
-            return io.NodeOutput(
-                [item["image"]], [item.get("mask")], [item["index"]], [context_id],
-                *[[item_values.get(name)] for name in VALUE_NAMES],
-            )
-
-        items = normalize_image_list_inputs(images, mask, values)
-        prompt = getattr(cls.hidden, "prompt", None)
-        start_id = str(getattr(cls.hidden, "unique_id", ""))
-        if not prompt or not start_id:
-            raise RuntimeError("Distributed Image List Start needs a submitted ComfyUI prompt context")
-        try:
-            worker_ids = json.loads(str(unwrap_single(image_list_enabled_worker_ids or "[]")))
-        except (TypeError, ValueError) as exc:
-            raise RuntimeError("Invalid preflighted image-list worker list") from exc
-        if not isinstance(worker_ids, list) or not all(isinstance(item, str) for item in worker_ids):
-            raise RuntimeError("Invalid preflighted image-list worker list")
-        from ..image_list.runtime import validate_image_list_regions
-        regions = validate_image_list_regions(prompt)
-        region = next((row for row in regions if row["start_id"] == start_id), None)
-        if region is None:
-            raise RuntimeError(f"No safe paired Image List End was found for Start node {start_id}")
-        from comfy_execution.utils import get_executing_context
-        executing = get_executing_context()
-        master_participates = unwrap_single(image_list_master_participates)
-        if master_participates is None:
-            from ..utils.config import load_config
-            master_participates = not bool((load_config().get("settings", {}) or {}).get("master_delegate_only", False))
-        if not isinstance(master_participates, bool):
-            raise ValueError("image_list_master_participates must be a boolean")
-        context_id = create_master_job(
-            start_id, region["end_id"], prompt, items, worker_ids,
-            str(unwrap_single(image_list_master_callback or "")),
-            image_list_workflow_meta if isinstance(image_list_workflow_meta, dict) else None,
-            prompt_id=executing.prompt_id if executing is not None else None,
-            master_participates=master_participates,
-        )
-        return io.NodeOutput(
-            [item["image"] for item in items],
-            [item.get("mask") for item in items],
-            [item["index"] for item in items],
-            [context_id for _ in items],
-            *[[item["values"].get(name) for item in items] for name in VALUE_NAMES],
-        )
-
-
-async def _consume_worker_item(token):
-    # Item state is local to each worker process; source nodes consume it exactly once.
-    from ..image_list.runtime import consume_worker_item
-    return consume_worker_item(token)
-
-
-class DistributedImageListEnd(io.ComfyNode):
-    """Collect branch results in source order without evaluating the branch on master."""
-
-    @classmethod
-    def define_schema(cls):
-        return io.Schema(
-            node_id="DistributedImageListEnd",
-            display_name="Distributed Image List End",
-            category="image/distributed",
-            description="Closes a paired image-list region and returns ordered image, mask, and text lists.",
-            inputs=[
-                io.Image.Input("images", lazy=True),
-                io.Mask.Input("masks", optional=True, lazy=True),
-                io.String.Input("texts", optional=True, lazy=True, force_input=True),
-                io.Custom("DISTRIBUTED_IMAGE_LIST_CONTEXT").Input("context"),
-            ],
-            outputs=[
-                io.Image.Output(display_name="image_list", is_output_list=True),
-                io.Mask.Output(display_name="masks", is_output_list=True),
-                io.String.Output(display_name="texts", is_output_list=True),
-            ],
-            is_input_list=True,
-            is_output_node=True,
-            accept_all_inputs=True,
-            not_idempotent=True,
-            hidden=[io.Hidden.prompt, io.Hidden.unique_id, io.Hidden.execution_list],
-        )
-
-    @classmethod
-    def check_lazy_status(cls, images=None, masks=None, texts=None, context=None,
-                          worker_result_token="", **kwargs):
-        worker_result_token = unwrap_single(worker_result_token)
-        if worker_result_token:
-            # Decide which optional sockets are connected from the actual worker prompt;
-            # absent optional inputs must not be demanded from Comfy's lazy executor.
-            prompt = getattr(cls.hidden, "prompt", {}) or {}
-            end_inputs = next((node.get("inputs", {}) for node in prompt.values()
-                               if isinstance(node, dict)
-                               and node.get("class_type") == "DistributedImageListEnd"
-                               and node.get("inputs", {}).get("worker_result_token") == worker_result_token), {})
-            required = []
-            for name, value in (("images", images), ("masks", masks), ("texts", texts)):
-                connected = name == "images" or name in end_inputs
-                missing = value is None or (isinstance(value, (list, tuple)) and all(v is None for v in value))
-                if connected and missing:
-                    required.append(name)
-            return required
-        # Master only needs Start's context. Never request branch outputs here.
-        if context is None or (isinstance(context, (tuple, list)) and all(item is None for item in context)):
-            return ["context"]
-        return []
-
-    @classmethod
-    def execute(cls, images=None, masks=None, texts=None, context=None,
-                worker_result_token="", worker_item_index=-1, master_callback_url="",
-                image_list_local_execution=False):
-        worker_result_token = unwrap_single(worker_result_token)
-        master_callback_url = unwrap_single(master_callback_url)
-        if worker_result_token:
-            context_values = [item for item in (context if isinstance(context, (list, tuple)) else [context])
-                              if isinstance(item, str) and item]
-            if not context_values or len(set(context_values)) != 1:
-                raise RuntimeError("Worker Image List End is missing its paired context")
-            result = normalize_worker_result(images, masks, texts)
-            if unwrap_single(image_list_local_execution):
-                from ..image_list.runtime import set_master_result
-                set_master_result(context_values[0], int(unwrap_single(worker_item_index)),
-                                  str(worker_result_token), result)
-            else:
-                payload = encode_result(result)
-                run_async_in_server_loop(
-                    submit_worker_result(str(master_callback_url), str(worker_result_token),
-                                         context_values[0], int(unwrap_single(worker_item_index)), payload), timeout=60,
-                )
-            return io.NodeOutput([result["image"]], [result.get("mask")], [result.get("text")])
-
-        context_values = [item for item in (context if isinstance(context, (list, tuple)) else [context])
-                          if isinstance(item, str) and item]
-        if not context_values or len(set(context_values)) != 1:
-            raise RuntimeError("Distributed Image List End requires the paired Start context wire")
-        context_id = context_values[0]
-        from ..image_list.runtime import ImageListCancelledError
-        try:
-            # The runtime owns the configurable deadline and bounded remote cleanup.
-            execution_list = getattr(cls.hidden, "execution_list", None)
-            # Local model outputs belong to this still-running outer prompt.
-            # Add them to its existing tracker; never start/end another tracker
-            # that could unpin a model still needed by the parent.
-            model_callback = getattr(execution_list, "output_link_callback", None)
-            result_lists = run_async_in_server_loop(run_image_list_job(context_id, model_callback=model_callback))
-        except ImageListCancelledError:
-            import comfy.model_management
-            raise comfy.model_management.InterruptProcessingException()
-        images_out, masks_out, texts_out = result_lists
-        if len(images_out) != len(masks_out) or len(images_out) != len(texts_out):
-            raise RuntimeError("Image-list collector received mismatched image, mask, or text result counts")
-        return io.NodeOutput(images_out, masks_out, texts_out)
-
-
-async def submit_worker_result(callback_url, item_token, context_id, item_index, result):
-    if not callback_url.startswith(("http://", "https://")):
-        raise RuntimeError("Worker image-list callback URL is missing or invalid")
-    import aiohttp
-    from ..utils.network import get_client_session
-    session = await get_client_session()
-    payload = {"context_id": context_id, "token": item_token,
-               "item_index": item_index, "result": result}
-    async with session.post(
-        f"{callback_url.rstrip('/')}/distributed/image-list/result",
-        data=encode_request_payload(payload),
-        headers={"Content-Type": "application/json"},
-        timeout=aiohttp.ClientTimeout(total=60),
-    ) as response:
-        if response.status >= 400:
-            body = (await response.text())[:500]
-            raise RuntimeError(f"Master rejected image-list result (HTTP {response.status}): {body}")

nodes/v3.py

@@ -1,215 +0,0 @@
-"""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
-from .image_list import DistributedImageListEnd, DistributedImageListStart
-
-
-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, DistributedImageListStart, DistributedImageListEnd]

pyproject.toml

@@ -1,7 +1,7 @@
 [project]
 name = "ComfyUI-Distributed"
 description = "ComfyUI extension that enables multi-GPU processing locally, remotely and in the cloud"
-version = "1.5.0"
+version = "1.4.7"
 license = {file = "LICENSE"}
 dependencies  = []
 

runtime/__init__.py

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

runtime/bootstrap.py

@@ -1,38 +0,0 @@
-"""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()
-    from ..image_list import runtime as image_list_runtime
-    from ..image_list.lifecycle import install_lifecycle
-    install_lifecycle(server.PromptServer.instance, image_list_runtime)
-
-    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/test_config_routes.py

@@ -38,10 +38,17 @@ def _load_config_routes_module():
     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")]
+    utils_pkg.__path__ = []
     sys.modules[f"{package_name}.utils"] = utils_pkg
 
-    import aiohttp  # Load the real dependency; isolate the fake response at the tested module.
+    created_aiohttp_stub = False
+    if "aiohttp" not in sys.modules:
+        created_aiohttp_stub = True
+        aiohttp_module = types.ModuleType("aiohttp")
+        aiohttp_module.web = types.SimpleNamespace(
+            json_response=lambda payload, status=200: _FakeResponse(payload, status=status)
+        )
+        sys.modules["aiohttp"] = aiohttp_module
 
     class _Routes:
         def get(self, _path):
@@ -89,14 +96,8 @@ def _load_config_routes_module():
     assert spec is not None and spec.loader is not None
     spec.loader.exec_module(module)
 
-    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()
+    if created_aiohttp_stub:
+        sys.modules.pop("aiohttp", None)
 
     return module
 

tests/api/test_distributed_queue.py

@@ -4,8 +4,6 @@ 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
@@ -28,17 +26,6 @@ class _FakeRequest:
         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"
@@ -59,7 +46,6 @@ def _load_job_routes_module():
     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
@@ -219,10 +205,6 @@ def _load_job_routes_module():
     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)
@@ -412,144 +394,5 @@ class JobCompleteAudioPayloadTests(unittest.IsolatedAsyncioTestCase):
             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_media_sync.py

@@ -164,11 +164,26 @@ 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)
@@ -243,121 +258,5 @@ class RewritePromptMediaInputsTests(unittest.TestCase):
         self.assertEqual(prompt["1"]["inputs"]["audio"], "song.wav")
 
 
-class AnnotatedMediaUploadTests(unittest.IsolatedAsyncioTestCase):
-    async def test_output_temp_and_input_uploads_do_not_collide(self):
-        from unittest.mock import AsyncMock, patch
-
-        class Response:
-            status = 200
-            def __init__(self, data):
-                self.data = data
-            async def __aenter__(self):
-                return self
-            async def __aexit__(self, *_args):
-                return None
-            async def json(self):
-                return self.data
-            def raise_for_status(self):
-                pass
-
-        class Form:
-            def __init__(self):
-                self.fields = {}
-            def add_field(self, key, value, **kwargs):
-                self.fields[key] = value
-                if key == 'image':
-                    self.fields['name'] = kwargs['filename']
-
-        checked = []
-        def post(url, *, json=None, data=None, **kwargs):
-            if url.endswith('/check_file'):
-                checked.append(json['filename'])
-                return Response({'exists': False})
-            self.assertEqual(data.fields['type'], 'input')
-            self.assertEqual(data.fields['image'], b'real file bytes')
-            return Response({'name': data.fields['name'], 'subfolder': data.fields['subfolder']})
-
-        refs = ['shared/ref.png', 'shared/ref.png [output]', 'shared/ref.png [temp]']
-        prompt = {str(i): {'inputs': {'image': ref}} for i, ref in enumerate(refs)}
-        self.assertEqual(ms._find_media_references(prompt), refs)
-        with patch.object(ms, 'get_client_session', AsyncMock(return_value=types.SimpleNamespace(post=post))), \
-             patch.object(ms, 'aiohttp', types.SimpleNamespace(FormData=Form, ClientTimeout=lambda **kw: kw)):
-            paths = [await ms._upload_media_to_worker({}, ref, b'real file bytes', 'fixture-hash', 'image/png')
-                     for ref in refs]
-        expected = ['shared/ref.png', 'distributed-media/output/fixture-hash/ref.png',
-                    'distributed-media/temp/fixture-hash/ref.png']
-        self.assertEqual(checked, expected)
-        self.assertEqual(paths, [(True, path) for path in expected])
-        ms._rewrite_prompt_media_inputs(prompt, dict(zip(refs, expected)))
-        self.assertEqual([node['inputs']['image'] for node in prompt.values()], expected)
-
-
 if __name__ == "__main__":
     unittest.main()
-
-
-# ---------------------------------------------------------------------------
-# 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_transfer_integration.py

@@ -1,219 +0,0 @@
-"""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_worker_routes.py

@@ -106,7 +106,7 @@ def _load_worker_routes_module():
     sys.modules[f"{package_name}.utils"] = utils_pkg
 
     workers_pkg = types.ModuleType(f"{package_name}.workers")
-    workers_pkg.__path__ = [str(module_path.parents[1] / "workers")]
+    workers_pkg.__path__ = []
     workers_pkg.get_worker_manager = lambda: _DummyWorkerManager()
     sys.modules[f"{package_name}.workers"] = workers_pkg
 
@@ -213,7 +213,6 @@ 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")
@@ -257,8 +256,6 @@ def _load_worker_routes_module():
     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))
 
     if created_aiohttp_stub:
         sys.modules.pop("aiohttp", None)
@@ -272,38 +269,6 @@ 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/fixtures/v1_node_contracts.json

@@ -1,691 +0,0 @@
-{
-  "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/image_list_network_check.py

@@ -1,475 +0,0 @@
-"""Model-free multi-process acceptance using real ComfyUI executors and loopback HTTP.
-
-Only the transport test facade (/prompt, /queue, /history, /interrupt) is synthetic.
-Start/End, graph extraction, scheduling, payload codecs and callbacks are production code.
-No managed ComfyUI service, user workflow, model, or remote worker is touched.
-"""
-import asyncio
-from collections import Counter
-import importlib
-import json
-import os
-from pathlib import Path
-import sys
-import tempfile
-import time
-import uuid
-from unittest.mock import patch
-
-ROOT = Path(__file__).resolve().parents[2]
-COMFY_ROOT = Path(sys.argv[1]).resolve()
-WORKER = '--worker' in sys.argv
-DELAY = float(sys.argv[3]) if WORKER else 0
-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
-import comfy.model_management
-import execution
-import nodes
-import server
-import torch
-from aiohttp import web
-import aiohttp
-
-torch.set_num_threads(1)
-SOURCE_CALLS, BRANCH_CALLS, CAPTURED, MODEL_CALLS = [], [], [], []
-ALLOW_MASTER = False
-COUNT = 7
-
-
-def source_rows():
-    return [{'image': torch.full((1, 5 + i, 7, 3 if i % 2 == 0 else 4), i / 16, dtype=torch.float32),
-             'reference_image': torch.full((1, 2, 3, 3), (i + 1) / 16),
-             'mask': torch.full((1, 5 + i, 7), (i + 1) / 16),
-             'text': f'item-{i}', 'strength': (i + 1) / 16, 'seed': 100 + i, 'index': i}
-            for i in range(COUNT)]
-
-
-class TestImageSource:
-    @classmethod
-    def INPUT_TYPES(cls):
-        return {'required': {}}
-    RETURN_TYPES = ('IMAGE', 'IMAGE', 'MASK', 'STRING', 'FLOAT', 'INT')
-    OUTPUT_IS_LIST = (True,) * 6
-    FUNCTION = 'run'
-    def run(self):
-        assert not WORKER, 'upstream image source leaked into worker prompt'
-        SOURCE_CALLS.append(True)
-        rows = source_rows()
-        return tuple([row[key] for row in rows] for key in
-                     ('image', 'reference_image', 'mask', 'text', 'strength', 'seed'))
-
-
-class TestSharedFlag:
-    @classmethod
-    def INPUT_TYPES(cls):
-        return {'required': {}}
-    RETURN_TYPES = ('BOOLEAN',)
-    FUNCTION = 'run'
-    def run(self):
-        assert not WORKER, 'shared value source leaked into worker prompt'
-        return (True,)
-
-
-class TestCachedModel:
-    @classmethod
-    def INPUT_TYPES(cls):
-        return {'required': {}}
-    RETURN_TYPES = ('MODEL',)
-    FUNCTION = 'run'
-    def run(self):
-        MODEL_CALLS.append(True)
-        return (object(),)
-
-
-class TestImageOperation:
-    @classmethod
-    def INPUT_TYPES(cls):
-        return {'required': {'image': ('IMAGE',), 'reference_image': ('IMAGE',), 'mask': ('MASK',),
-                             'text': ('STRING', {'forceInput': True}),
-                             'strength': ('FLOAT', {'forceInput': True}),
-                             'seed': ('INT', {'forceInput': True}), 'index': ('INT', {'forceInput': True}),
-                             'flag': ('BOOLEAN', {'forceInput': True}), 'model': ('MODEL',)}}
-    RETURN_TYPES = ('IMAGE', 'MASK', 'STRING')
-    FUNCTION = 'run'
-    def run(self, image, reference_image, mask, text, strength, seed, index, flag, model):
-        assert WORKER or ALLOW_MASTER, 'per-image branch executed in orchestrator-only mode'
-        BRANCH_CALLS.append(index)
-        # Sleep in short chunks so the real Comfy interrupt signal can be tested.
-        until = time.monotonic() + DELAY
-        while time.monotonic() < until:
-            comfy.model_management.throw_exception_if_processing_interrupted()
-            time.sleep(0.01)
-        assert seed == 100 + index
-        assert flag is True
-        assert text == f'item-{index}'
-        assert strength == (index + 1) / 16
-        assert torch.all(reference_image == strength)
-        return image + strength, mask * 0.5, f'{text}:seed-{seed}'
-
-
-class TestCapture:
-    @classmethod
-    def INPUT_TYPES(cls):
-        return {'required': {'images': ('IMAGE',), 'masks': ('MASK',), 'texts': ('STRING', {'forceInput': True})}}
-    INPUT_IS_LIST = True
-    OUTPUT_NODE = True
-    RETURN_TYPES = ()
-    FUNCTION = 'run'
-    def run(self, images, masks, texts):
-        CAPTURED.append((images, masks, texts))
-        return ()
-
-
-async def setup():
-    assets = default_asset_manager()
-    ps = server.PromptServer(asyncio.get_running_loop(), assets)
-    assert await nodes.load_custom_node(str(ROOT))
-    module = sys.modules[str(ROOT).replace('.', '_x_')]
-    nodes.NODE_CLASS_MAPPINGS.update(TestImageSource=TestImageSource,
-                                   TestImageOperation=TestImageOperation, TestCapture=TestCapture,
-                                   TestCachedModel=TestCachedModel, TestSharedFlag=TestSharedFlag)
-    # Match the CORS-enabled direct-worker launch contract for canvas probes.
-    app = web.Application(middlewares=[server.create_cors_middleware('*')])
-    # Register the actual feature routes, not test replacements.
-    app.add_routes([route for route in ps.routes if route.path.startswith('/distributed/image-list/')])
-    return ps, assets, module, app
-
-
-async def listen(app):
-    runner = web.AppRunner(app)
-    await runner.setup()
-    site = web.TCPSite(runner, '127.0.0.1', 0)
-    await site.start()
-    return runner, site._server.sockets[0].getsockname()[1]
-
-
-async def worker_main():
-    ps, assets, module, app = await setup()
-    pending, running, history = [], {}, {}
-    wake = asyncio.Event()
-    interrupted_ids = []
-    async def prompt(request):
-        data = await request.json()
-        prompt_id = str(uuid.uuid4())
-        graph = data['prompt']
-        valid = await execution.validate_prompt(prompt_id, graph, None)
-        if not valid[0]:
-            return web.json_response({'error': valid[1], 'node_errors': valid[3]}, status=400)
-        pending.append((prompt_id, graph, valid[2]))
-        wake.set()
-        return web.json_response({'prompt_id': prompt_id})
-    async def queue(request):
-        if request.method == 'POST':
-            deleted = (await request.json()).get('delete', [])
-            pending[:] = [entry for entry in pending if entry[0] not in deleted]
-            return web.json_response({})
-        return web.json_response({'queue_running': [[0, key] for key in running],
-                                  'queue_pending': [[0, entry[0]] for entry in pending]})
-    async def get_history(request):
-        pid = request.match_info['prompt_id']
-        return web.json_response({pid: history[pid]} if pid in history else {})
-    async def interrupt(request):
-        pid = (await request.json())['prompt_id']
-        interrupted_ids.append(pid)
-        if pid in running:
-            comfy.model_management.interrupt_current_processing(True)
-        return web.json_response({})
-    async def stats(request):
-        return web.json_response({'calls': BRANCH_CALLS, 'running': list(running),
-                                  'pending': [e[0] for e in pending], 'interrupted': interrupted_ids,
-                                  'history': history})
-    async def consume():
-        executor = execution.PromptExecutor(ps, cache_args={'ram': 0, 'ram_inactive': 0}, asset_manager=assets)
-        while True:
-            await wake.wait()
-            while pending:
-                pid, graph, outputs = pending.pop(0)
-                running[pid] = True
-                await asyncio.to_thread(executor.execute, graph, pid, {}, outputs)
-                history[pid] = {'status': {'completed': True,
-                    'status_str': 'success' if executor.success else 'error',
-                    'messages': str(executor.status_messages)}}
-                running.pop(pid, None)
-            wake.clear()
-    async def probe(_request):
-        return web.json_response({'exec_info': {'queue_remaining': len(pending) + len(running)}})
-    app.router.add_get('/prompt', probe)
-    app.router.add_post('/prompt', prompt)
-    app.router.add_get('/queue', queue)
-    app.router.add_post('/queue', queue)
-    app.router.add_get('/history/{prompt_id}', get_history)
-    app.router.add_post('/interrupt', interrupt)
-    app.router.add_get('/test-stats', stats)
-    runner, port = await listen(app)
-    consumer = asyncio.create_task(consume())
-    print(f'READY:{port}', flush=True)
-    try:
-        await asyncio.Event().wait()
-    finally:
-        consumer.cancel()
-        await asyncio.gather(consumer, return_exceptions=True)
-        await runner.cleanup()
-
-
-def graph(callback, worker_ids):
-    return {
-        '1': {'class_type': 'TestImageSource', 'inputs': {}},
-        # Generic Autogrow value slots carry per-item reference images, text,
-        # strength and seed (one value per image), plus a shared boolean.
-        '2': {'class_type': 'DistributedImageListStart', 'inputs': {
-            'images': ['1', 0], 'mask': ['1', 2],
-            'values.value_1': ['1', 1], 'values.value_2': ['1', 3],
-            'values.value_3': ['1', 4], 'values.value_4': ['1', 5], 'values.value_5': ['7', 0],
-            'image_list_enabled_worker_ids': json.dumps(worker_ids), 'image_list_master_callback': callback}},
-        '3': {'class_type': 'TestImageOperation', 'inputs': {
-            key: ['2', index] for key, index in [('image', 0), ('mask', 1), ('index', 2),
-                ('reference_image', 4), ('text', 5), ('strength', 6), ('seed', 7), ('flag', 8)]}
-              | {'model': ['6', 0]}},
-        '4': {'class_type': 'DistributedImageListEnd', 'inputs': {
-            'images': ['3', 0], 'masks': ['3', 1], 'texts': ['3', 2], 'context': ['2', 3]}},
-        '5': {'class_type': 'TestCapture', 'inputs': {'images': ['4', 0], 'masks': ['4', 1], 'texts': ['4', 2]}},
-        '6': {'class_type': 'TestCachedModel', 'inputs': {}},
-        '7': {'class_type': 'TestSharedFlag', 'inputs': {}},
-    }
-
-
-async def coordinator():
-    global ALLOW_MASTER
-    ps, assets, module, app = await setup()
-    runtime = importlib.import_module(module.__name__ + '.image_list.runtime')
-    config_mod = importlib.import_module(module.__name__ + '.utils.config')
-    orchestration = importlib.import_module(module.__name__ + '.api.queue_orchestration')
-    runner, port = await listen(app)
-    ps.port = port  # Match the listener used by local-worker callback resolution.
-    processes, stderr_files = [], []
-    # Old worker capability refusal is read-only and precedes input evaluation.
-    legacy_app = web.Application()
-    legacy_requests = []
-    async def legacy(request):
-        legacy_requests.append(request.method)
-        return web.json_response({'error': 'not found'}, status=404)
-    legacy_app.router.add_route('*', '/{path:.*}', legacy)
-    legacy_runner, legacy_port = await listen(legacy_app)
-    dispatch_module = importlib.import_module(module.__name__ + '.api.orchestration.dispatch')
-    try:
-        try:
-            await dispatch_module.preflight_image_list_workers([
-                {'id': 'legacy', 'type': 'local', 'host': '127.0.0.1', 'port': legacy_port}])
-            raise AssertionError('legacy worker accepted')
-        except RuntimeError as exc:
-            assert 'upgrade' in str(exc), str(exc)
-        assert legacy_requests == ['GET'] and not SOURCE_CALLS
-    finally:
-        await legacy_runner.cleanup()
-    print('IMAGE_LIST_PREFLIGHT_OK legacy worker rejected before materialization or writes')
-    scratch = Path(os.environ.get('TMPDIR', Path.home() / '.hermes/cache/scratch'))
-    scratch.mkdir(parents=True, exist_ok=True)
-    try:
-        with tempfile.TemporaryDirectory(prefix='image-list-network-', dir=scratch) as tmp:
-            workers = []
-            for name, delay in [('fast', 0.02), ('slow', 0.5)]:
-                err = open(Path(tmp) / f'{name}.log', 'w+')
-                stderr_files.append(err)
-                proc = await asyncio.create_subprocess_exec(sys.executable, str(Path(__file__).resolve()),
-                    str(COMFY_ROOT), '--worker', str(delay), stdout=asyncio.subprocess.PIPE, stderr=err,
-                    env={**os.environ, 'TMPDIR': tmp})
-                processes.append(proc)
-                async def ready():
-                    while True:
-                        line = await proc.stdout.readline()
-                        if not line:
-                            err.seek(0)
-                            raise AssertionError(f'Worker exited: {err.read()}')
-                        if line.startswith(b'READY:'):
-                            return int(line.split(b':')[1])
-                worker_port = await asyncio.wait_for(ready(), 60)
-                workers.append({'id': name, 'type': 'local', 'host': '127.0.0.1',
-                                'port': worker_port, 'enabled': True})
-            config = {'workers': workers, 'settings': {'websocket_orchestration': False, 'master_delegate_only': True},
-                      'master': {'host': f'http://127.0.0.1:{port}'}}
-            prompt = graph(f'http://127.0.0.1:{port}', [w['id'] for w in workers])
-            valid = await execution.validate_prompt('network-acceptance', prompt, None)
-            assert valid[0], valid
-            executor = execution.PromptExecutor(ps, cache_args={'ram': 0, 'ram_inactive': 0}, asset_manager=assets)
-            # Exercise production queue preparation; only master enqueue is redirected
-            # to this isolated executor, not the managed server's live queue.
-            prompt['2']['inputs'].pop('image_list_enabled_worker_ids')
-            prompt['2']['inputs'].pop('image_list_master_callback')
-            async def enqueue(prepared, workflow, client_id, **kwargs):
-                checked = await execution.validate_prompt('network-acceptance', prepared, None)
-                assert checked[0], checked
-                await asyncio.to_thread(executor.execute, prepared, 'network-acceptance', {}, checked[2])
-                return {'prompt_id': 'network-acceptance', 'number': 1, 'node_errors': {}}
-            with patch.object(config_mod, 'load_config', return_value=config), \
-                 patch.object(orchestration, 'load_config', return_value=config), \
-                 patch.object(orchestration, 'queue_prompt_payload', enqueue), \
-                 patch.object(config_mod, 'get_worker_timeout_seconds', return_value=15):
-                queued = await asyncio.wait_for(orchestration.orchestrate_distributed_execution(
-                    prompt, None, 'test', [w['id'] for w in workers]), 45)
-                assert queued[0] == 'network-acceptance' and queued[2] == 2, queued
-                assert not ps.distributed_pending_jobs, 'image-list created unused Collector queues'
-            if not executor.success:
-                for err in stderr_files:
-                    err.flush(); err.seek(0); print(err.read(), file=sys.stderr)
-            assert executor.success, executor.status_messages
-            assert len(SOURCE_CALLS) == 1 and not BRANCH_CALLS
-            assert len(CAPTURED) == 1
-            images, masks, texts = CAPTURED[0]
-            assert len(images) == len(masks) == len(texts) == COUNT
-            for i, row in enumerate(source_rows()):
-                assert torch.equal(images[i], row['image'] + row['strength']), (i, images[i])
-                assert torch.equal(masks[i], row['mask'] * 0.5), i
-                assert texts[i] == f'item-{i}:seed-{100 + i}', texts
-            async with aiohttp.ClientSession() as session:
-                stats = []
-                for worker in workers:
-                    async with session.get(f"http://127.0.0.1:{worker['port']}/test-stats") as response:
-                        stats.append(await response.json())
-                counts = [len(row['calls']) for row in stats]
-                assert counts[0] > counts[1] > 0, counts
-                assert Counter(i for row in stats for i in row['calls']) == Counter(range(COUNT))
-            assert not runtime._MASTER_JOBS
-            # Three sequential regions on one worker, including optional outputs
-            # left unconnected. Region N+1 must not race the prior worker End.
-            chained = {'1': {'class_type': 'EmptyImage', 'inputs': {
-                'height': 8, 'width': 8, 'batch_size': 2, 'color': 0}}}
-            previous = '1'
-            for start, end in [('2', '3'), ('4', '5'), ('6', '7')]:
-                chained[start] = {'class_type': 'DistributedImageListStart', 'inputs': {
-                    'images': [previous, 0]}}
-                chained[end] = {'class_type': 'DistributedImageListEnd', 'inputs': {
-                    'images': [start, 0], 'context': [start, 3]}}
-                previous = end
-            chained['8'] = {'class_type': 'TestCapture', 'inputs': {
-                'images': ['7', 0], 'masks': ['7', 1], 'texts': ['7', 2]}}
-            with patch.object(config_mod, 'load_config', return_value=config), \
-                 patch.object(orchestration, 'load_config', return_value=config), \
-                 patch.object(orchestration, 'queue_prompt_payload', enqueue), \
-                 patch.object(config_mod, 'get_worker_timeout_seconds', return_value=15):
-                await asyncio.wait_for(orchestration.orchestrate_distributed_execution(
-                    chained, None, 'test', ['slow']), 45)
-            assert executor.success, executor.status_messages
-            assert len(CAPTURED) == 2 and len(CAPTURED[1][0]) == 2
-            assert all(torch.equal(image, torch.zeros((1, 8, 8, 3))) for image in CAPTURED[1][0])
-            assert CAPTURED[1][1:] == ([None, None], [None, None])
-            assert not runtime._MASTER_JOBS
-            print('IMAGE_LIST_SEQUENTIAL_OK three regions; one worker; optional outputs omitted')
-            # Same public queue path, now with master participation enabled. The
-            # parent queue stays occupied while its local slot processes items.
-            ALLOW_MASTER = True
-            config['settings']['master_delegate_only'] = False
-            participation_counts = {}
-            for mode, selected in [('mixed', [w['id'] for w in workers]), ('master-only', [])]:
-                executor = execution.PromptExecutor(ps, cache_args={'ram': 0, 'ram_inactive': 0}, asset_manager=assets)
-                before_sources, before_local, before_models = len(SOURCE_CALLS), len(BRANCH_CALLS), len(MODEL_CALLS)
-                async with aiohttp.ClientSession() as session:
-                    before_worker_calls = []
-                    for worker in workers:
-                        async with session.get(f"http://127.0.0.1:{worker['port']}/test-stats") as response:
-                            before_worker_calls.append(len((await response.json())['calls']))
-                with patch.object(config_mod, 'load_config', return_value=config), \
-                     patch.object(orchestration, 'load_config', return_value=config), \
-                     patch.object(orchestration, 'queue_prompt_payload', enqueue), \
-                     patch.object(config_mod, 'get_worker_timeout_seconds', return_value=15):
-                    queued = await asyncio.wait_for(orchestration.orchestrate_distributed_execution(
-                        prompt, None, 'test', selected), 45)
-                assert executor.success, (mode, executor.status_messages)
-                assert queued[2] == len(selected)
-                assert len(SOURCE_CALLS) == before_sources + 1, 'source was reevaluated per local item'
-                local_calls = BRANCH_CALLS[before_local:]
-                assert local_calls, f'{mode}: master did not process any items'
-                assert len(MODEL_CALLS) == before_models + 1, 'local model dependency did not stay cached'
-                remote_calls = []
-                async with aiohttp.ClientSession() as session:
-                    for worker, offset in zip(workers, before_worker_calls):
-                        async with session.get(f"http://127.0.0.1:{worker['port']}/test-stats") as response:
-                            current = await response.json()
-                        remote_calls.extend(current['calls'][offset:])
-                        assert not current['running'] and not current['pending'], current
-                if mode == 'mixed':
-                    assert remote_calls, 'mixed execution failed to use direct workers'
-                else:
-                    assert not remote_calls and len(local_calls) == COUNT
-                assert Counter(local_calls + remote_calls) == Counter(range(COUNT)), (local_calls, remote_calls)
-                images, masks, texts = CAPTURED[-1]
-                assert len(images) == len(masks) == len(texts) == COUNT
-                for i, row in enumerate(source_rows()):
-                    assert torch.equal(images[i], row['image'] + row['strength']), (mode, i)
-                    assert torch.equal(masks[i], row['mask'] * 0.5), (mode, i)
-                    assert texts[i] == f'item-{i}:seed-{100 + i}', (mode, texts)
-                assert not runtime._MASTER_JOBS and not runtime._WORKER_ITEMS
-                participation_counts[mode] = {'master': len(local_calls), 'workers': len(remote_calls)}
-            print('IMAGE_LIST_MASTER_OK ' + json.dumps(participation_counts))
-            ALLOW_MASTER = False
-            config['settings']['master_delegate_only'] = True
-            # Real Comfy user-interrupt flag while a real worker executor is active.
-            slow = workers[1]
-            async with aiohttp.ClientSession() as session:
-                async def slow_stats():
-                    async with session.get(f"http://127.0.0.1:{slow['port']}/test-stats") as response:
-                        return await response.json()
-                async def until(predicate):
-                    for _ in range(150):
-                        current = await slow_stats()
-                        if predicate(current):
-                            return current
-                        await asyncio.sleep(0.02)
-                    raise AssertionError('worker state transition timed out')
-                before = await until(lambda state: not state['running'] and not state['pending'])
-                context = runtime.create_master_job('2', '4', prompt, source_rows(), ['slow'],
-                                                  f'http://127.0.0.1:{port}')
-                with patch.object(config_mod, 'load_config', return_value=config), \
-                     patch.object(config_mod, 'get_worker_timeout_seconds', return_value=15):
-                    task = asyncio.create_task(runtime.run_image_list_job(context))
-                    active = await until(lambda state: bool(state['running']))
-                    owned_id = active['running'][0]
-                    comfy.model_management.interrupt_current_processing(True)
-                    try:
-                        await asyncio.wait_for(task, 5)
-                        raise AssertionError('interrupt was ignored')
-                    except runtime.ImageListCancelledError:
-                        pass
-                    finally:
-                        comfy.model_management.interrupt_current_processing(False)
-                after = await until(lambda state: not state['running'] and not state['pending'])
-                assert after['interrupted'] == [owned_id], after
-                assert len(after['calls']) <= len(before['calls']) + 1
-                assert set(before['history']) <= set(after['history'])
-                assert not runtime._MASTER_JOBS
-            print('IMAGE_LIST_CANCEL_OK owned prompt stopped; no further dispatch; context removed')
-            # A worker exception before End must fail, not hang or return a partial list.
-            bad_rows = source_rows()
-            bad_rows[0]['seed'] = -1
-            context = runtime.create_master_job('2', '4', prompt, bad_rows, ['fast'],
-                                              f'http://127.0.0.1:{port}')
-            with patch.object(config_mod, 'load_config', return_value=config), \
-                 patch.object(config_mod, 'get_worker_timeout_seconds', return_value=15):
-                try:
-                    await asyncio.wait_for(runtime.run_image_list_job(context), 5)
-                    raise AssertionError('worker exception was ignored')
-                except RuntimeError as exc:
-                    assert 'failed image-list item' in str(exc), str(exc)
-            assert not runtime._MASTER_JOBS
-            print('IMAGE_LIST_FAILURE_OK worker exception surfaced; context removed')
-            print('IMAGE_LIST_NETWORK_OK ' + json.dumps({'workers': 2, 'items': COUNT,
-                'assignments': dict(zip(['fast', 'slow'], counts)), 'ordered_exact_outputs': True,
-                'master_branch_calls': len(BRANCH_CALLS), 'source_calls': len(SOURCE_CALLS)}))
-    finally:
-        for proc in processes:
-            if proc.returncode is None:
-                proc.terminate()
-        await asyncio.gather(*(proc.wait() for proc in processes))
-        for err in stderr_files:
-            err.close()
-        await runner.cleanup()
-        network = importlib.import_module(module.__name__ + '.utils.network')
-        session = await network.get_client_session()
-        await session.close()
-
-
-if __name__ == '__main__':
-    asyncio.run(worker_main() if WORKER else coordinator())

tests/helpers/v3_runtime_check.py

@@ -1,480 +0,0 @@
-"""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', 'DistributedImageListStart', 'DistributedImageListEnd'}
-
-
-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 check_image_list_lazy_executor(mapping, module, prompt_server, asset_manager):
-    """Execute real native V3 list/lazy flow; master must skip the processing branch."""
-    branch_calls = []
-    class ImageListProbe:
-        @classmethod
-        def INPUT_TYPES(cls):
-            return {'required': {'image': ('IMAGE',)}}
-        RETURN_TYPES = ('IMAGE',)
-        FUNCTION = 'run'
-        CATEGORY = 'tests'
-        def run(self, image):
-            branch_calls.append(image)
-            return (image,)
-
-    nodes.NODE_CLASS_MAPPINGS['ImageListTestProbe'] = ImageListProbe
-    runtime = sys.modules[module.__name__ + '.image_list.runtime']
-    original_contexts = set(runtime._MASTER_JOBS)
-    image_list_nodes = sys.modules[module.__name__ + '.nodes.image_list']
-    workflow_meta = {'nodes': [{'id': 2, 'type': 'DistributedImageListStart'}], 'extra': {'acceptance': True}}
-    async def fake_worker_results(_context_id, *, model_callback=None):
-        state = runtime.get_master_job(_context_id)
-        assert state.prompt_id == 'image-list-lazy', state.prompt_id
-        assert state.workflow_meta == workflow_meta, state.workflow_meta
-        assert len(prompt_server.app.cleanup_ctx) >= 1, 'lifecycle cleanup was not registered'
-        return (
-            [torch.full((1, 8, 8, 3), value, dtype=torch.float32) for value in (0.125, 0.5, 0.875)],
-            [None, None, None], [None, None, None],
-        )
-
-    graph = {
-        '1': {'class_type': 'EmptyImage', 'inputs': {'height': 8, 'width': 8, 'batch_size': 3, 'color': 0}},
-        '2': {'class_type': 'DistributedImageListStart', 'inputs': {
-            'images': ['1', 0], 'image_list_enabled_worker_ids': '["test-worker"]',
-            'image_list_master_callback': 'http://127.0.0.1:1',
-            'image_list_workflow_meta': workflow_meta,
-        }},
-        '3': {'class_type': 'ImageListTestProbe', 'inputs': {'image': ['2', 0]}},
-        '4': {'class_type': 'DistributedImageListEnd', 'inputs': {
-            'images': ['3', 0], 'context': ['2', 3],
-        }},
-        '5': {'class_type': 'PreviewImage', 'inputs': {'images': ['4', 0]}},
-    }
-    valid = await execution.validate_prompt('image-list-lazy', graph, None)
-    assert valid[0], valid
-    validator = runtime.validate_image_list_regions
-    assert 'texts' not in graph['4']['inputs']
-    assert len(validator(graph)) == 1
-    connected_text_graph = json.loads(json.dumps(graph))
-    connected_text_graph['6'] = {
-        'class_type': 'ImageListTextProbe', 'inputs': {'text': ['2', 4]},
-    }
-    connected_text_graph['4']['inputs']['texts'] = ['6', 0]
-    assert len(validator(connected_text_graph)) == 1
-    for literal_text in ('', 'per-item text'):
-        literal_text_graph = json.loads(json.dumps(graph))
-        literal_text_graph['4']['inputs']['texts'] = literal_text
-        try:
-            validator(literal_text_graph)
-        except ValueError:
-            pass
-        else:
-            raise AssertionError(f'End accepted unconnected literal text: {literal_text!r}')
-    print('IMAGE_LIST_END_TEXT_INPUT_OK native forceInput; unconnected omitted; linked accepted; literals rejected')
-    import folder_paths
-    scratch = Path(os.environ.get('TMPDIR', Path.home() / '.hermes/cache/scratch'))
-    try:
-        # Resolve genuine Comfy output/temp annotations against the pinned core,
-        # not only the controlled folder_paths seam in unit tests.
-        media = sys.modules[module.__name__ + '.api.orchestration.media_sync']
-        with tempfile.TemporaryDirectory(prefix='media-roots-', dir=scratch) as media_root:
-            for kind in ('input', 'output', 'temp'):
-                source = Path(media_root) / kind / 'shared/ref.png'
-                source.parent.mkdir(parents=True)
-                source.write_bytes(kind.encode())
-            with patch.object(folder_paths, 'input_directory', str(Path(media_root) / 'input')), \
-                 patch.object(folder_paths, 'output_directory', str(Path(media_root) / 'output')), \
-                 patch.object(folder_paths, 'temp_directory', str(Path(media_root) / 'temp')):
-                for kind in ('output', 'temp'):
-                    refs = media._find_media_references({'1': {'inputs': {'image': f'shared/ref.png [{kind}]'}}})
-                    assert len(refs) == 1
-                    assert media._load_media_file_sync(refs[0])[0] == kind.encode()
-        with tempfile.TemporaryDirectory(prefix='image-list-lazy-', dir=scratch) as temp:
-            with patch.object(folder_paths, 'temp_directory', temp), \
-                 patch.object(image_list_nodes, 'run_image_list_job', fake_worker_results):
-                executor = execution.PromptExecutor(
-                    prompt_server, cache_args={'ram': 0, 'ram_inactive': 0}, asset_manager=asset_manager)
-                await asyncio.to_thread(executor.execute, graph, 'image-list-lazy', {}, valid[2])
-                assert executor.success, executor.status_messages
-                assert branch_calls == [], 'master evaluated the lazy per-image branch'
-                records = executor.history_result['outputs']['5']['images']
-                assert len(records) == 3, records
-                for record in records:
-                    preview = Path(temp) / record.get('subfolder', '') / record['filename']
-                    with Image.open(preview) as image:
-                        assert image.size == (8, 8) and image.mode == 'RGB'
-    finally:
-        nodes.NODE_CLASS_MAPPINGS.pop('ImageListTestProbe', None)
-        for context_id in set(runtime._MASTER_JOBS) - original_contexts:
-            runtime.remove_master_job(context_id)
-    print('IMAGE_LIST_LAZY_EXECUTION_OK real PromptExecutor; branch skipped on master; 3 previews verified')
-
-
-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) == 11
-    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'))
-    start_schema = mapping['DistributedImageListStart'].GET_SCHEMA()
-    end_schema = mapping['DistributedImageListEnd'].GET_SCHEMA()
-    assert start_schema.node_id == 'DistributedImageListStart'
-    assert start_schema.display_name == 'Distributed Image List Start'
-    assert end_schema.node_id == 'DistributedImageListEnd'
-    assert end_schema.display_name == 'Distributed Image List End'
-    texts_input = next(input for input in end_schema.inputs if input.id == 'texts')
-    texts_schema = texts_input.as_dict()
-    assert texts_schema.get('forceInput') is True, texts_schema
-    assert 'default' not in texts_schema, texts_schema
-    assert [out.io_type for out in start_schema.outputs] == [
-        'IMAGE', 'MASK', 'INT', 'DISTRIBUTED_IMAGE_LIST_CONTEXT'] + ['*'] * 10
-    assert [out.display_name for out in start_schema.outputs][4:] == [f'value_{i}' for i in range(1, 11)]
-    assert all(out.is_output_list for out in start_schema.outputs)
-    assert start_schema.is_input_list and start_schema.accept_all_inputs
-    start_info = json.loads(json.dumps(mapping['DistributedImageListStart'].GET_NODE_INFO_V1()))
-    assert list(start_info['input']['optional']) == ['images', 'mask', 'values'], start_info['input']
-    values_spec = start_info['input']['optional']['values']
-    assert values_spec[0] == 'COMFY_AUTOGROW_V3', values_spec
-    template = values_spec[1]['template']
-    assert template['names'] == [f'value_{i}' for i in range(1, 11)] and template['min'] == 0, template
-    assert template['input']['required']['value'][0] == 'IMAGE,MASK,STRING,INT,FLOAT,BOOLEAN', template
-    # Autogrow + INPUT_IS_LIST through the real parser/list mapper, values
-    # still split 1-or-N, unconnected slots ignored.
-    # get_input_data must accept Autogrow names; constants are wrapped once.
-    parsed, missing, hidden = execution.get_input_data(
-        {'values.value_1': 'shared', 'values.value_3': 4},
-        mapping['DistributedImageListStart'], unique_id='autogrow-probe', extra_data={})
-    assert not missing and parsed['values.value_1'] == ['shared'], parsed
-    # Linked inputs arrive as the upstream output list (one element per value).
-    prepared = {'images': [torch.zeros((3, 2, 2, 3))], 'values.value_1': ['shared'],
-                'values.value_3': [4, 5, 6]}
-    seen = {}
-    runtime_mod = sys.modules[module.__name__ + '.image_list.runtime']
-    def capture(_images, _mask, values):
-        seen['values'] = values
-        raise RuntimeError('captured')
-    nodes_mod = sys.modules[module.__name__ + '.nodes.image_list']
-    with patch.object(nodes_mod, 'normalize_image_list_inputs', capture):
-        try:
-            await asyncio.to_thread(lambda: asyncio.run(execution._async_map_node_over_list(
-                'autogrow-probe', 'autogrow-probe', mapping['DistributedImageListStart'], prepared,
-                'execute', v3_data=hidden)))
-        except RuntimeError as exc:
-            assert 'captured' in str(exc), exc
-    assert seen['values'] == {'value_1': ['shared'], 'value_3': [4, 5, 6]}, seen
-    rows = runtime_mod.normalize_image_list_inputs(prepared['images'], None, seen['values'])
-    assert [row['values'] for row in rows] == [{'value_1': 'shared', 'value_3': n} for n in (4, 5, 6)]
-    print('IMAGE_LIST_AUTOGROW_OK real parser; value_N slots shared/per-item')
-    assert [out.io_type for out in end_schema.outputs] == ['IMAGE', 'MASK', 'STRING']
-    assert all(out.is_output_list for out in end_schema.outputs)
-    assert end_schema.is_input_list and end_schema.is_output_node
-    print('IMAGE_LIST_SCHEMA_OK exact IDs, names, typed list outputs')
-    print('SCHEMA_PARITY_OK', len(mapping))
-    check_saved_workflows(mapping)
-    await check_execution(mapping, module, prompt_server, asset_manager)
-    await check_image_list_lazy_executor(mapping, module, prompt_server, asset_manager)
-    print('V3_ACCEPTANCE_OK')
-
-
-asyncio.run(main())

tests/test_collector_image_transport.py

@@ -0,0 +1,225 @@
+"""Real PNG/tensor regressions; only server/config/HTTP plumbing is stubbed."""
+import asyncio
+import importlib.util
+import io
+import base64
+import sys
+import types
+from pathlib import Path
+from unittest.mock import patch
+
+import numpy as np
+import pytest
+import torch
+from aiohttp import web
+from PIL import Image
+
+
+ROOT = Path(__file__).resolve().parents[1]
+
+
+def _load_file(name, path):
+    spec = importlib.util.spec_from_file_location(name, path)
+    module = importlib.util.module_from_spec(spec)
+    sys.modules[name] = module
+    spec.loader.exec_module(module)
+    return module
+
+
+@pytest.fixture
+def media_modules():
+    # Existing harnesses provide ComfyUI/config doubles. Restore their module
+    # replacements so this file can coexist with the rest of the test suite.
+    with patch.dict(sys.modules):
+        image_utils = _load_file("real_image_utils", ROOT / "utils/image.py")
+        collector_tests = _load_file("collector_harness", ROOT / "tests/test_collector_list_inputs.py")
+        collector = collector_tests._load_collector_module()
+        sys.modules["aiohttp"].web = web
+        route_tests = _load_file("route_harness", ROOT / "tests/api/test_distributed_queue.py")
+        routes = route_tests._load_job_routes_module()
+        for module in (collector, routes):
+            module.pil_to_tensor = image_utils.pil_to_tensor
+            module.ensure_contiguous = image_utils.ensure_contiguous
+        collector.tensor_to_pil = image_utils.tensor_to_pil
+        routes.prompt_server = collector.prompt_server
+        yield collector, routes
+
+
+def _png_payload(image):
+    buffer = io.BytesIO()
+    image.save(buffer, format="PNG")
+    return "data:image/png;base64," + base64.b64encode(buffer.getvalue()).decode("ascii")
+
+
+@pytest.mark.parametrize("mode", ["RGB", "RGBA", "L", "LA", "palette", "palette_alpha"])
+def test_png_decoder_preserves_alpha_and_normalizes_opaque_modes(media_modules, mode):
+    _, routes = media_modules
+    if mode.startswith("palette"):
+        image = Image.new("P", (12, 8), 0)
+        image.putpalette([16, 32, 64] + [0] * 765)
+        if mode == "palette_alpha":
+            image.info["transparency"] = 0
+    else:
+        color = {"RGB": (16, 32, 64), "RGBA": (16, 32, 64, 96), "L": 32, "LA": (32, 96)}[mode]
+        image = Image.new(mode, (12, 8), color)
+    alpha = mode in {"RGBA", "LA", "palette_alpha"}
+    expected = np.array(image.convert("RGBA" if alpha else "RGB"), dtype=np.float32) / 255
+    decoded = routes._decode_canonical_png_tensor(_png_payload(image))
+    assert decoded.shape == (1, 8, 12, 4 if alpha else 3)
+    assert decoded.is_contiguous()
+    np.testing.assert_array_equal(decoded[0].numpy(), expected)
+
+
+async def _master_collect(module, local, worker_images, delegate=False):
+    module.prompt_server.distributed_jobs_lock = asyncio.Lock()
+    queue = asyncio.Queue()
+    for index, tensor in enumerate(worker_images):
+        queue.put_nowait({"worker_id": "worker-a", "image_index": index,
+                          "tensor": tensor, "is_last": index == len(worker_images) - 1})
+    module.prompt_server.distributed_pending_jobs = {"image-test": queue}
+    return await module.DistributedCollectorNode().execute(
+        local, None, multi_job_id="image-test", enabled_worker_ids='["worker-a"]',
+        delegate_only=delegate,
+    )
+
+
+@pytest.mark.parametrize("local_channels,worker_channels", [(3, 3), (4, 4), (3, 4), (4, 3), (None, 4)])
+def test_worker_png_callback_and_collector_preserve_pixels_and_alpha(
+    media_modules, local_channels, worker_channels
+):
+    module, routes = media_modules
+    local = None if local_channels is None else torch.full((1, 8, 12, local_channels), 0.5)
+    worker = torch.full((2, 8, 12, worker_channels), 0.25)
+    if worker_channels == 4:
+        worker[..., 3] = torch.linspace(0, 1, 12).reshape(1, 1, 12)
+    responses = []
+
+    async def scenario():
+        module.prompt_server.distributed_jobs_lock = asyncio.Lock()
+        module.prompt_server.distributed_pending_jobs = {"image-test": asyncio.Queue()}
+
+        class Response:
+            def __init__(self, payload):
+                self.payload = payload
+
+            async def __aenter__(self):
+                async def request_json():
+                    return self.payload
+                response = await routes.job_complete_endpoint(types.SimpleNamespace(json=request_json))
+                responses.append(response.status)
+                assert response.status == 200, response.text
+                return self
+
+            async def __aexit__(self, *_args):
+                return False
+
+            def raise_for_status(self):
+                pass
+
+        class Session:
+            def post(self, _url, json, timeout):
+                return Response(json)
+
+        async def get_session():
+            return Session()
+
+        module.get_client_session = get_session
+        module.encode_audio_payload = lambda _audio: None
+        collector = module.DistributedCollectorNode()
+        await collector.send_batch_to_master(worker, None, "image-test", "http://master", "worker-a")
+        result, _ = await collector.execute(
+            local, None, multi_job_id="image-test", enabled_worker_ids='["worker-a"]',
+            delegate_only=local is None,
+        )
+        assert "image-test" not in module.prompt_server.distributed_pending_jobs
+        return result
+
+    result = asyncio.run(scenario())
+    channels = max(local_channels or 3, worker_channels)
+    offset = 0 if local is None else 1
+    assert responses == [200, 200]
+    assert result.shape == (offset + 2, 8, 12, channels)
+    assert result.is_contiguous()
+    if local is not None:
+        assert torch.equal(result[:1, ..., :local_channels], local)
+        if local_channels == 3 and channels == 4:
+            assert torch.all(result[:1, ..., 3] == 1)
+    assert torch.allclose(result[offset:, ..., :worker_channels], worker, atol=1 / 255)
+    if worker_channels == 3 and channels == 4:
+        assert torch.all(result[offset:, ..., 3] == 1)
+
+
+def test_collector_mixed_rgb_rgba_list_promotes_rgb_to_opaque_rgba(media_modules):
+    module, _ = media_modules
+    rgb = torch.zeros(1, 8, 12, 3)
+    rgba = torch.full((1, 8, 12, 4), 0.5)
+    result, _ = module.DistributedCollectorNode().run(images=[rgb, rgba])
+    assert result.shape == (2, 8, 12, 4)
+    assert torch.equal(result[:1, ..., :3], rgb)
+    assert torch.all(result[:1, ..., 3] == 1)
+    assert torch.equal(result[1:], rgba)
+
+
+@pytest.mark.parametrize("other_channels", [1, 3, 4])
+def test_collector_preserves_grayscale_or_promotes_it_for_color_batches(media_modules, other_channels):
+    module, _ = media_modules
+    gray = torch.full((1, 8, 12, 1), 0.25)
+    other = torch.full((1, 8, 12, other_channels), 0.5)
+    result, _ = module.DistributedCollectorNode().run(images=[gray, other])
+    assert result.shape == (2, 8, 12, other_channels)
+    assert torch.all(result[:1, ..., :min(other_channels, 3)] == 0.25)
+    if other_channels == 4:
+        assert torch.all(result[:1, ..., 3] == 1)
+    assert torch.equal(result[1:], other)
+
+
+def test_mixed_channel_merge_keeps_configured_worker_and_image_order(media_modules):
+    module, _ = media_modules
+    collector = module.DistributedCollectorNode()
+    master = torch.zeros(1, 8, 12, 3)
+    worker_a = torch.full((1, 8, 12, 4), 0.25)
+    worker_b = torch.full((1, 8, 12, 3), 0.5)
+    unexpected = torch.full((1, 8, 12, 3), 0.75)
+    result = collector._reorder_and_combine_tensors(
+        {"worker-a": {0: worker_a}, "worker-b": {2: worker_b, 0: master},
+         "unexpected": {0: unexpected}},
+        ["worker-b", "worker-a"], 1, master, False, master,
+    )
+    assert result.shape == (5, 8, 12, 4)
+    assert result[:, 0, 0, 0].tolist() == [0, 0, 0.5, 0.25, 0.75]
+    assert result[:, 0, 0, 3].tolist() == [1, 1, 1, 0.25, 1]
+
+
+def test_collector_spatial_mismatch_raises_with_shapes_not_successful_fallback(media_modules):
+    module, _ = media_modules
+    local = torch.zeros(1, 8, 12, 3)
+    worker = torch.zeros(1, 8, 13, 3)
+    with pytest.raises(ValueError, match="spatial dimensions") as error:
+        asyncio.run(_master_collect(module, local, [worker]))
+    message = str(error.value)
+    assert "master image 0" in message
+    assert "worker 'worker-a' image 0" in message
+    assert "(1, 8, 12, 3)" in message and "(1, 8, 13, 3)" in message
+    assert "image-test" not in module.prompt_server.distributed_pending_jobs
+
+
+@pytest.mark.parametrize("shape", [(1, 8, 12), (1, 8, 12, 5)])
+def test_collector_rejects_invalid_worker_image_shape(media_modules, shape):
+    module, _ = media_modules
+    with pytest.raises(ValueError, match="worker 'worker-a' image 0"):
+        asyncio.run(_master_collect(module, None, [torch.zeros(shape)], delegate=True))
+
+
+def test_delegate_master_ignores_empty_placeholder_shape(media_modules):
+    module, _ = media_modules
+    placeholder = torch.empty(0, 64, 64, 3)
+    worker = torch.full((1, 8, 12, 4), 0.5)
+    result, _ = asyncio.run(_master_collect(module, placeholder, [worker], delegate=True))
+    assert torch.equal(result, worker)
+
+
+def test_empty_image_results_raise_instead_of_reaching_save_image(media_modules):
+    module, _ = media_modules
+    with pytest.raises(ValueError, match="empty image batch"):
+        asyncio.run(_master_collect(module, torch.empty(0, 64, 64, 3),
+                                    [torch.empty(0, 8, 12, 3)], delegate=True))

tests/test_collector_list_inputs.py

@@ -7,17 +7,6 @@ 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"
@@ -37,7 +26,6 @@ def _load_collector_module():
     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):
@@ -98,8 +86,6 @@ def _load_collector_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")
@@ -376,255 +362,3 @@ def test_delegate_only_audio_collects_worker_audio_without_placeholder_image():
     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

@@ -18,7 +18,7 @@ def _load_config_module():
             del sys.modules[mod_name]
 
     root_pkg = types.ModuleType(package_name)
-    root_pkg.__path__ = [str(module_path.parent)]
+    root_pkg.__path__ = []
     sys.modules[package_name] = root_pkg
 
     logging_module = types.ModuleType(f"{package_name}.logging")

tests/test_dispatch_selection.py

@@ -55,13 +55,7 @@ def _load_dispatch_module():
             def __init__(self, total=None):
                 self.total = total
 
-        class _ClientError(Exception):
-            pass
-
-        class _ClientConnectorError(_ClientError):
-            pass
-
-        class _ClientConnectionError(_ClientError):
+        class _ClientConnectorError(Exception):
             pass
 
         class _WSMsgType:
@@ -82,9 +76,7 @@ def _load_dispatch_module():
                 self.closed = True
 
         aiohttp_module.ClientTimeout = _ClientTimeout
-        setattr(aiohttp_module, "ClientError", _ClientError)
         aiohttp_module.ClientConnectorError = _ClientConnectorError
-        setattr(aiohttp_module, "ClientConnectionError", _ClientConnectionError)
         aiohttp_module.WSMsgType = _WSMsgType
         aiohttp_module.TCPConnector = _TCPConnector
         aiohttp_module.ClientSession = _ClientSession
@@ -108,75 +100,6 @@ def _load_dispatch_module():
 dispatch = _load_dispatch_module()
 
 
-def _load_prompt_transform_module():
-    module_path = Path(__file__).resolve().parents[1] / "api" / "orchestration" / "prompt_transform.py"
-    package_name = "prompt_transform_worker_selection_testpkg"
-    for suffix in ("", ".api", ".api.orchestration", ".utils"):
-        package = types.ModuleType(package_name + suffix)
-        package.__path__ = []
-        sys.modules[package_name + suffix] = package
-    logging_module = types.ModuleType(f"{package_name}.utils.logging")
-    setattr(logging_module, "debug_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
-
-
-prompt_transform = _load_prompt_transform_module()
-
-
-class _FakeResponse:
-    def __init__(self, status=200, payload=None, error=None):
-        self.status = status
-        self.payload = payload
-        self.error = error
-
-    async def __aenter__(self):
-        if self.error is not None:
-            raise self.error
-        return self
-
-    async def __aexit__(self, *_args):
-        return None
-
-    async def json(self):
-        return self.payload
-
-
-class _FakeSession:
-    def __init__(self, responses):
-        self.responses = responses
-
-    def get(self, url, timeout=None):
-        del timeout
-        worker_id = url.removeprefix("http://").split("/", 1)[0]
-        endpoint = "capabilities" if url.endswith("/capabilities") else "queue"
-        response = self.responses[(worker_id, endpoint)]
-        if isinstance(response, BaseException):
-            return _FakeResponse(error=response)
-        status, payload = response
-        return _FakeResponse(status=status, payload=payload)
-
-
-def _valid_capability():
-    return {
-        "supported": True,
-        "protocol_version": 2,
-        "node_ids": ["DistributedImageListStart", "DistributedImageListEnd"],
-    }
-
-
-def _valid_queue(running=None, pending=None):
-    return {"queue_running": running or [], "queue_pending": pending or []}
-
-
 class DispatchSelectionTests(unittest.IsolatedAsyncioTestCase):
     async def test_select_active_workers_filters_offline(self):
         workers = [
@@ -326,194 +249,6 @@ class DispatchSelectionTests(unittest.IsolatedAsyncioTestCase):
 
         self.assertIsNone(selected)
 
-    async def test_image_list_preflight_skips_unreachable_worker_and_keeps_healthy_worker(self):
-        workers = [{"id": "healthy"}, {"id": "offline"}]
-        responses = {
-            (worker_id, "capabilities"): (200, _valid_capability())
-            for worker_id in ("healthy", "offline")
-        }
-        responses.update({
-            (worker_id, "queue"): (200, _valid_queue())
-            for worker_id in ("healthy", "offline")
-        })
-        responses[("offline", "capabilities")] = dispatch.aiohttp.ClientConnectionError("connection refused")
-        session = _FakeSession(responses)
-
-        async def fake_session():
-            return session
-
-        with patch.object(dispatch, "get_client_session", side_effect=fake_session), patch.object(
-            dispatch, "build_worker_url", side_effect=lambda worker, path: f"http://{worker['id']}{path}"
-        ), patch.object(dispatch, "log") as log_mock:
-            result = await dispatch.preflight_image_list_workers(workers)
-
-        self.assertEqual(result, {"healthy": True})
-        self.assertTrue(any("offline" in str(call) for call in log_mock.call_args_list))
-
-    async def test_image_list_preflight_skips_timeout_and_keeps_healthy_worker(self):
-        workers = [{"id": "healthy"}, {"id": "slow"}]
-        responses: dict = {
-            (worker_id, "capabilities"): (200, _valid_capability())
-            for worker_id in ("healthy", "slow")
-        }
-        responses.update({
-            (worker_id, "queue"): (200, _valid_queue())
-            for worker_id in ("healthy", "slow")
-        })
-        responses[("slow", "queue")] = asyncio.TimeoutError()
-        session = _FakeSession(responses)
-
-        async def fake_session():
-            return session
-
-        with patch.object(dispatch, "get_client_session", side_effect=fake_session), patch.object(
-            dispatch, "build_worker_url", side_effect=lambda worker, path: f"http://{worker['id']}{path}"
-        ), patch.object(dispatch, "log") as log_mock:
-            result = await dispatch.preflight_image_list_workers(workers)
-
-        self.assertEqual(result, {"healthy": True})
-        self.assertTrue(any("slow" in str(call) for call in log_mock.call_args_list))
-
-    async def test_image_list_preflight_excludes_busy_workers(self):
-        workers = [{"id": "idle"}, {"id": "busy"}]
-        responses = {
-            (worker_id, "capabilities"): (200, _valid_capability())
-            for worker_id in ("idle", "busy")
-        }
-        responses.update({
-            ("idle", "queue"): (200, _valid_queue()),
-            ("busy", "queue"): (200, _valid_queue(running=[{"prompt": "existing"}])),
-        })
-        session = _FakeSession(responses)
-
-        async def fake_session():
-            return session
-
-        with patch.object(dispatch, "get_client_session", side_effect=fake_session), patch.object(
-            dispatch, "build_worker_url", side_effect=lambda worker, path: f"http://{worker['id']}{path}"
-        ):
-            result = await dispatch.preflight_image_list_workers(workers)
-
-        self.assertEqual(result, {"idle": True})
-
-    async def test_image_list_preflight_reports_no_idle_workers_when_all_are_unavailable(self):
-        workers = [{"id": "offline-a"}, {"id": "offline-b"}]
-        responses = {
-            (worker_id, "capabilities"): dispatch.aiohttp.ClientConnectionError("unreachable")
-            for worker_id in ("offline-a", "offline-b")
-        }
-        session = _FakeSession(responses)
-
-        async def fake_session():
-            return session
-
-        with patch.object(dispatch, "get_client_session", side_effect=fake_session), patch.object(
-            dispatch, "build_worker_url", side_effect=lambda worker, path: f"http://{worker['id']}{path}"
-        ):
-            self.assertEqual(
-                await dispatch.preflight_image_list_workers(workers, require_idle=False),
-                {},
-            )
-            with self.assertRaisesRegex(RuntimeError, "No selected feature-compatible direct worker is idle"):
-                await dispatch.preflight_image_list_workers(workers, require_idle=True)
-
-    async def test_image_list_preflight_fails_closed_for_malformed_capabilities(self):
-        malformed = {"supported": True, "protocol_version": 2, "node_ids": None}
-        session = _FakeSession({("worker", "capabilities"): (200, malformed)})
-
-        async def fake_session():
-            return session
-
-        with patch.object(dispatch, "get_client_session", side_effect=fake_session), patch.object(
-            dispatch, "build_worker_url", side_effect=lambda worker, path: f"http://{worker['id']}{path}"
-        ):
-            with self.assertRaisesRegex(RuntimeError, "malformed.*capabilit"):
-                await dispatch.preflight_image_list_workers([{"id": "worker"}])
-
-    async def test_image_list_preflight_fails_closed_for_incompatible_protocol(self):
-        incompatible = _valid_capability()
-        incompatible["protocol_version"] = 1
-        session = _FakeSession({("worker", "capabilities"): (200, incompatible)})
-
-        async def fake_session():
-            return session
-
-        with patch.object(dispatch, "get_client_session", side_effect=fake_session), patch.object(
-            dispatch, "build_worker_url", side_effect=lambda worker, path: f"http://{worker['id']}{path}"
-        ):
-            with self.assertRaisesRegex(RuntimeError, "incompatible"):
-                await dispatch.preflight_image_list_workers([{"id": "worker"}])
-
-    async def test_image_list_preflight_fails_closed_for_missing_required_node(self):
-        incompatible = _valid_capability()
-        incompatible["node_ids"] = ["DistributedImageListStart"]
-        session = _FakeSession({("worker", "capabilities"): (200, incompatible)})
-
-        async def fake_session():
-            return session
-
-        with patch.object(dispatch, "get_client_session", side_effect=fake_session), patch.object(
-            dispatch, "build_worker_url", side_effect=lambda worker, path: f"http://{worker['id']}{path}"
-        ):
-            with self.assertRaisesRegex(RuntimeError, "incompatible"):
-                await dispatch.preflight_image_list_workers([{"id": "worker"}])
-
-    async def test_image_list_preflight_fails_closed_for_malformed_queue(self):
-        responses = {
-            ("worker", "capabilities"): (200, _valid_capability()),
-            ("worker", "queue"): (200, {"queue_running": [], "queue_pending": None}),
-        }
-        session = _FakeSession(responses)
-
-        async def fake_session():
-            return session
-
-        with patch.object(dispatch, "get_client_session", side_effect=fake_session), patch.object(
-            dispatch, "build_worker_url", side_effect=lambda worker, path: f"http://{worker['id']}{path}"
-        ):
-            with self.assertRaisesRegex(RuntimeError, "malformed /queue"):
-                await dispatch.preflight_image_list_workers([{"id": "worker"}])
-
-    async def test_image_list_preflight_does_not_hide_unexpected_errors(self):
-        session = _FakeSession({("worker", "capabilities"): ValueError("programming error")})
-
-        async def fake_session():
-            return session
-
-        with patch.object(dispatch, "get_client_session", side_effect=fake_session), patch.object(
-            dispatch, "build_worker_url", side_effect=lambda worker, path: f"http://{worker['id']}{path}"
-        ):
-            with self.assertRaisesRegex(ValueError, "programming error"):
-                await dispatch.preflight_image_list_workers([{"id": "worker"}])
-
-
-class ImageListDirectWorkerResolutionTests(unittest.TestCase):
-    def setUp(self):
-        self.config = {"workers": [
-            {"id": "enabled", "enabled": True, "host": "enabled-host", "port": 8188},
-            {"id": "disabled", "enabled": False, "host": "disabled-host", "port": 8189},
-            {"id": "serverless", "enabled": True, "type": "serverless"},
-            {"id": "other", "enabled": True, "host": "other-host", "port": 8190},
-        ]}
-
-    def test_implicit_selection_uses_only_enabled_direct_workers(self):
-        workers = prompt_transform._resolve_direct_workers_for_image_list(self.config, None)
-        self.assertEqual([worker["id"] for worker in workers], ["enabled", "other"])
-
-    def test_explicit_selection_includes_exact_disabled_direct_worker(self):
-        workers = prompt_transform._resolve_direct_workers_for_image_list(self.config, ["disabled"])
-        self.assertEqual([worker["id"] for worker in workers], ["disabled"])
-
-    def test_explicit_empty_selection_returns_no_workers(self):
-        self.assertEqual(prompt_transform._resolve_direct_workers_for_image_list(self.config, []), [])
-
-    def test_unknown_explicit_id_raises_instead_of_falling_back(self):
-        with self.assertRaisesRegex(RuntimeError, "missing"):
-            prompt_transform._resolve_direct_workers_for_image_list(self.config, ["enabled", "missing"])
-
-    def test_explicit_serverless_selection_is_excluded(self):
-        self.assertEqual(prompt_transform._resolve_direct_workers_for_image_list(self.config, ["serverless"]), [])
-
 
 if __name__ == "__main__":
     unittest.main()

tests/test_image_list.py

@@ -1,333 +0,0 @@
-"""Regression tests for typed dynamic image-list normalization and graph scheduling."""
-import asyncio
-import unittest
-
-import torch
-
-from image_list.runtime import (
-    END_CLASS,
-    IMAGE_OUT,
-    CONTEXT_OUT,
-    START_CLASS,
-    VALUE_OUT_BASE,
-    ImageListCancelledError,
-    build_worker_prompt,
-    create_master_job,
-    decode_item,
-    decode_result,
-    dynamic_schedule,
-    encode_item,
-    encode_result,
-    normalize_image_list_inputs,
-    normalize_worker_result,
-    remove_master_job,
-    set_master_result,
-    validate_image_list_regions,
-)
-
-
-def start_node(image_link):
-    return {"class_type": START_CLASS, "inputs": {"images": image_link}}
-
-
-def graph_one_region():
-    return {
-        "1": {"class_type": "CheckpointLoaderSimple", "inputs": {}},
-        "2": {"class_type": "EmptyImage", "inputs": {"height": 8, "width": 8, "batch_size": 2}},
-        "3": start_node(["2", 0]),
-        "4": {"class_type": "GenericFrameOperation", "inputs": {
-            "image": ["3", IMAGE_OUT], "model": ["1", 0], "seed": ["3", VALUE_OUT_BASE],
-        }},
-        "5": {"class_type": END_CLASS, "inputs": {
-            "images": ["4", 0], "context": ["3", CONTEXT_OUT],
-        }},
-        "6": {"class_type": "SaveImage", "inputs": {"images": ["5", 0]}},
-    }
-
-
-def graph_two_sequential_regions():
-    return {
-        "1": {"class_type": "EmptyImage", "inputs": {"height": 4, "width": 4, "batch_size": 2}},
-        "2": start_node(["1", 0]),
-        "3": {"class_type": "FirstFrameOperation", "inputs": {"image": ["2", IMAGE_OUT]}},
-        "4": {"class_type": END_CLASS, "inputs": {"images": ["3", 0], "context": ["2", CONTEXT_OUT]}},
-        "5": start_node(["4", 0]),
-        "6": {"class_type": "SecondFrameOperation", "inputs": {"image": ["5", IMAGE_OUT]}},
-        "7": {"class_type": END_CLASS, "inputs": {"images": ["6", 0], "context": ["5", CONTEXT_OUT]}},
-        "8": {"class_type": "SaveImage", "inputs": {"images": ["7", 0]}},
-    }
-
-
-class ImageListNormalizationTests(unittest.TestCase):
-    def test_batch_frame_normalization_preserves_sizes_channels_and_float_values(self):
-        images = [
-            torch.tensor([[[[0.125, 0.25, 0.5]]]], dtype=torch.float64),
-            torch.arange(2 * 3 * 4, dtype=torch.float32).reshape(1, 2, 3, 4) / 17,
-        ]
-        refs = torch.ones((2, 5, 7, 3), dtype=torch.float16) * 0.333
-        masks = [torch.full((1, 2, 3), 0.125, dtype=torch.float64),
-                 torch.full((1, 8, 9), 0.875, dtype=torch.float32)]
-        rows = normalize_image_list_inputs(images, masks, {
-            "value_1": [refs[:1], refs[1:]], "value_2": ["first", "second"],
-            "value_3": [0.2], "value_4": [37],
-        })
-        self.assertEqual(len(rows), 2)
-        self.assertEqual([tuple(row["image"].shape) for row in rows], [(1, 1, 1, 3), (1, 2, 3, 4)])
-        self.assertEqual([row["image"].dtype for row in rows], [torch.float64, torch.float32])
-        self.assertTrue(torch.equal(rows[0]["values"]["value_1"], refs[:1]))
-        self.assertEqual(rows[1]["mask"].dtype, torch.float32)
-        self.assertEqual([row["index"] for row in rows], [0, 1])
-        self.assertEqual([row["values"]["value_2"] for row in rows], ["first", "second"])
-        self.assertEqual([row["values"]["value_3"] for row in rows], [0.2, 0.2])
-        self.assertEqual([row["values"]["value_4"] for row in rows], [37, 37])
-
-    def test_one_value_is_shared_and_n_values_pair_in_order(self):
-        image = torch.zeros((3, 2, 4, 3), dtype=torch.float32)
-        rows = normalize_image_list_inputs(image, values={
-            "value_1": ["same"], "value_2": [0.4], "value_3": [9], "value_4": [True]})
-        self.assertEqual([row["values"]["value_1"] for row in rows], ["same"] * 3)
-        self.assertEqual([row["values"]["value_2"] for row in rows], [0.4] * 3)
-        self.assertEqual([row["values"]["value_3"] for row in rows], [9] * 3)
-        self.assertEqual([row["values"]["value_4"] for row in rows], [True] * 3)
-        rows = normalize_image_list_inputs(image, values={"value_2": [1, 2, 3]})
-        self.assertEqual([row["values"] for row in rows], [{"value_2": 1}, {"value_2": 2}, {"value_2": 3}])
-        with self.assertRaisesRegex(ValueError, "value_1 has 2 values; connect exactly 1.*or 3"):
-            normalize_image_list_inputs(image, values={"value_1": ["one", "two"]})
-        with self.assertRaisesRegex(ValueError, "mask has 2 values"):
-            normalize_image_list_inputs(image, mask=torch.zeros((2, 2, 4)))
-        shared_mask = normalize_image_list_inputs(image, mask=torch.ones((1, 2, 4)))
-        self.assertTrue(all(torch.equal(row["mask"], torch.ones((1, 2, 4))) for row in shared_mask))
-
-    def test_value_batches_stay_whole_and_only_images_split(self):
-        images = torch.zeros((2, 2, 2, 3), dtype=torch.float32)
-        pair = torch.rand((2, 4, 4, 3))  # e.g. two references for one item
-        rows = normalize_image_list_inputs(images, values={"value_1": [pair]})
-        self.assertTrue(all(torch.equal(row["values"]["value_1"], pair) for row in rows))
-        per_item = [torch.rand((2, 4, 4, 3)), torch.rand((2, 4, 4, 3))]
-        rows = normalize_image_list_inputs(images, values={"value_1": per_item})
-        self.assertTrue(torch.equal(rows[1]["values"]["value_1"], per_item[1]))
-
-    def test_rejects_unsupported_values_and_counts(self):
-        images = torch.zeros((2, 2, 2, 3), dtype=torch.float32)
-        with self.assertRaisesRegex(ValueError, "value_1 has 3 values"):
-            normalize_image_list_inputs(images, values={"value_1": [1, 2, 3]})
-        with self.assertRaisesRegex(TypeError, "unsupported type dict"):
-            normalize_image_list_inputs(images, values={"value_1": [{"samples": 1}]})
-        with self.assertRaisesRegex(TypeError, "floating-point"):
-            normalize_image_list_inputs(images, values={"value_1": [torch.zeros(2, dtype=torch.int64)]})
-        with self.assertRaisesRegex(TypeError, "finite"):
-            normalize_image_list_inputs(images, values={"value_1": [float("nan")]})
-        with self.assertRaisesRegex(ValueError, "Unknown Start value slot"):
-            normalize_image_list_inputs(images, values={"value_11": [1]})
-        with self.assertRaisesRegex(ValueError, "at least one image"):
-            normalize_image_list_inputs(None)
-        # Unconnected Autogrow slots arrive as None and are ignored.
-        rows = normalize_image_list_inputs(images, values={"value_1": [None], "value_2": None})
-        self.assertEqual([row["values"] for row in rows], [{}, {}])
-
-    def test_typed_transport_round_trips_float_images_masks_and_scalar_values(self):
-        item = {"image": torch.arange(24, dtype=torch.float64).reshape(1, 2, 4, 3) / 7,
-                "mask": torch.full((1, 5, 6), 0.375, dtype=torch.float32),
-                "values": {"value_1": torch.ones((2, 1, 3, 4), dtype=torch.float16),
-                           "value_2": "safe text", "value_3": 0.625, "value_4": 42,
-                           "value_5": False},
-                "index": 3}
-        restored = decode_item(encode_item(item))
-        self.assertTrue(torch.equal(restored["image"], item["image"]))
-        self.assertEqual(restored["image"].dtype, torch.float64)
-        self.assertTrue(torch.equal(restored["values"]["value_1"], item["values"]["value_1"]))
-        self.assertEqual(restored["values"]["value_1"].dtype, torch.float16)
-        self.assertTrue(torch.equal(restored["mask"], item["mask"]))
-        self.assertEqual(restored["values"]["value_2"], "safe text")
-        self.assertEqual(restored["values"]["value_3"], 0.625)
-        self.assertIs(restored["values"]["value_4"], 42)
-        self.assertIs(restored["values"]["value_5"], False)
-        with self.assertRaisesRegex(ValueError, "unknown slot"):
-            decode_item({**encode_item(item), "values": {"strength": {"type": "FLOAT", "value": 1.0}}})
-        with self.assertRaisesRegex(ValueError, "unsupported type"):
-            decode_item({**encode_item(item), "values": {"value_1": {"type": "PATH", "value": "/x"}}})
-        with self.assertRaisesRegex(TypeError, "floating-point tensor"):
-            encode_item({**item, "image": torch.ones((1, 2, 4, 3), dtype=torch.int32)})
-        with self.assertRaisesRegex(ValueError, "Unsupported image-list transport type"):
-            from image_list.runtime import encode_value
-            encode_value("/untrusted/path", "PATH")
-
-    def test_result_transport_preserves_ragged_result_shapes(self):
-        result = normalize_worker_result(
-            [torch.ones((1, 3, 5, 4), dtype=torch.float32)],
-            [torch.ones((1, 3, 5), dtype=torch.float16)],
-            ["done"],
-        )
-        encoded = encode_result(result)
-        decoded = decode_result(encoded)
-        self.assertEqual(tuple(decoded["image"].shape), (1, 3, 5, 4))
-        self.assertEqual(decoded["image"].dtype, torch.float32)
-        self.assertEqual(decoded["mask"].dtype, torch.float16)
-        self.assertEqual(decoded["text"], "done")
-        with self.assertRaisesRegex(ValueError, "exactly one IMAGE"):
-            normalize_worker_result(torch.zeros((2, 3, 4, 3)))
-
-
-class ImageListGraphBoundaryTests(unittest.TestCase):
-    def test_prunes_upstream_item_source_but_keeps_external_model_loaders(self):
-        prompt = graph_one_region()
-        region = validate_image_list_regions(prompt)[0]
-        worker = build_worker_prompt(prompt, region, "token_12345678901234567890123456789012",
-                                     "ctx", 1, "http://master.invalid")
-        self.assertEqual(set(worker), {"1", "3", "4", "5"})
-        self.assertEqual(worker["3"]["inputs"]["worker_item_token"], "token_12345678901234567890123456789012")
-        self.assertNotIn("images", worker["3"]["inputs"])
-
-    def test_worker_prompt_drops_per_item_value_links(self):
-        prompt = graph_one_region()
-        prompt["8"] = {"class_type": "StringList", "inputs": {}}
-        prompt["3"]["inputs"].update({"values.value_1": ["8", 0], "values.value_2": ["8", 0]})
-        import sys, types
-        from unittest.mock import patch
-        registry = {"CheckpointLoaderSimple": type("Loader", (), {"RETURN_TYPES": ("MODEL",)})}
-        with patch.dict(sys.modules, {"nodes": types.SimpleNamespace(NODE_CLASS_MAPPINGS=registry)}):
-            region = validate_image_list_regions(prompt)[0]
-        worker = build_worker_prompt(prompt, region, "token_12345678901234567890123456789012", "ctx", 0, "http://m")
-        self.assertFalse([name for name in worker["3"]["inputs"] if name.startswith("values.")])
-        self.assertNotIn("8", worker)
-        self.assertEqual(worker["4"]["inputs"]["model"], ["1", 0])
-        self.assertNotIn("6", worker)
-
-    def test_rejects_upstream_processing_and_leaf_image_bypasses(self):
-        # Both a leaf source and a generated batch must cross Start exactly once.
-        for generated in (False, True):
-            with self.subTest(generated=generated):
-                prompt = graph_one_region()
-                if generated:
-                    prompt["2"] = {"class_type": "VAEDecode", "inputs": {"samples": ["7", 0]}}
-                    prompt["7"] = {"class_type": "KSampler", "inputs": {"model": ["1", 0]}}
-                prompt["4"]["inputs"]["reference"] = ["2", 0]
-                with self.assertRaisesRegex(ValueError, "bypasses Start.*2"):
-                    validate_image_list_regions(prompt)
-
-    def test_shared_typed_loader_and_leaf_constant_remain_valid(self):
-        import sys
-        import types
-        from unittest.mock import patch
-        prompt = graph_one_region()
-        prompt["7"] = {"class_type": "SeedConstant", "inputs": {"value": 9}}
-        prompt["2"]["inputs"].update({"model": ["1", 0], "seed": ["7", 0]})
-        prompt["4"]["inputs"]["shared_seed"] = ["7", 0]
-        registry = {
-            "CheckpointLoaderSimple": type("Loader", (), {"RETURN_TYPES": ("MODEL", "CLIP", "VAE")}),
-            "SeedConstant": type("Constant", (), {"RETURN_TYPES": ("INT",)}),
-        }
-        with patch.dict(sys.modules, {"nodes": types.SimpleNamespace(NODE_CLASS_MAPPINGS=registry)}):
-            region = validate_image_list_regions(prompt)[0]
-        worker = build_worker_prompt(prompt, region, "token", "ctx", 0, "http://master")
-        self.assertNotIn("2", worker)
-        self.assertTrue({"1", "7"}.issubset(worker))
-
-    def test_unknown_or_list_producing_shared_dependency_is_not_exempt(self):
-        import sys
-        import types
-        from unittest.mock import patch
-        for node_class in (None, type("ListLoader", (), {
-                "RETURN_TYPES": ("MODEL",), "OUTPUT_IS_LIST": (True,)})):
-            prompt = graph_one_region()
-            prompt["2"]["inputs"]["model"] = ["1", 0]
-            registry = {} if node_class is None else {"CheckpointLoaderSimple": node_class}
-            with patch.dict(sys.modules, {"nodes": types.SimpleNamespace(NODE_CLASS_MAPPINGS=registry)}):
-                with self.assertRaisesRegex(ValueError, "bypasses Start.*1"):
-                    validate_image_list_regions(prompt)
-
-    def test_two_sequential_regions_are_valid_and_isolated(self):
-        regions = validate_image_list_regions(graph_two_sequential_regions())
-        self.assertEqual([(r["start_id"], r["end_id"]) for r in regions], [("2", "4"), ("5", "7")])
-        p = graph_two_sequential_regions()
-        second = build_worker_prompt(p, regions[1], "item_token_123456789012345678901234", "ctx", 0,
-                                     "http://master.invalid")
-        self.assertEqual(set(second), {"5", "6", "7"})
-
-    def test_rejects_missing_context_pair_nested_region_and_cross_boundary_branch(self):
-        prompt = graph_one_region()
-        prompt["5"]["inputs"]["context"] = ["2", 0]
-        with self.assertRaisesRegex(ValueError, "typed context output"):
-            validate_image_list_regions(prompt)
-        nested = graph_one_region()
-        nested["7"] = start_node(["4", 0])
-        nested["8"] = {"class_type": "NestedOperation", "inputs": {"image": ["7", IMAGE_OUT]}}
-        nested["5"]["inputs"]["images"] = ["8", 0]
-        nested["9"] = {"class_type": END_CLASS, "inputs": {"images": ["8", 0], "context": ["7", CONTEXT_OUT]}}
-        with self.assertRaisesRegex(ValueError, "Nested/ambiguous"):
-            validate_image_list_regions(nested)
-        outside = graph_one_region()
-        outside["7"] = {"class_type": "PreviewImage", "inputs": {"images": ["3", IMAGE_OUT]}}
-        with self.assertRaisesRegex(ValueError, "cross-boundary"):
-            validate_image_list_regions(outside)
-
-    def test_rejects_duplicate_or_missing_results(self):
-        prompt = graph_one_region()
-        prompt["5"]["inputs"].pop("images")
-        with self.assertRaisesRegex(ValueError, "images must be connected"):
-            validate_image_list_regions(prompt)
-
-
-class DynamicImageSchedulingTests(unittest.IsolatedAsyncioTestCase):
-    async def test_fast_workers_get_more_items_and_order_is_stable(self):
-        assignments = {"fast": [], "slow": []}
-        async def dispatch(worker, index):
-            assignments[worker].append(index)
-            delay = 0.002 if worker == "fast" else 0.025
-            task = asyncio.create_task(asyncio.sleep(delay, result=(index, f"result-{index}")))
-            return {"index": index, "task": task}
-        async def wait_next(inflight):
-            done, _ = await asyncio.wait([h["task"] for h in inflight.values()], return_when=asyncio.FIRST_COMPLETED)
-            task = next(iter(done))
-            worker = next(w for w, handle in inflight.items() if handle["task"] is task)
-            index, result = task.result()
-            return worker, index, result
-        async def cleanup(_worker, handle):
-            handle["task"].cancel()
-        ordered = await dynamic_schedule(12, ["fast", "slow"], dispatch, wait_next, cleanup)
-        self.assertEqual(ordered, [f"result-{i}" for i in range(12)])
-        self.assertGreater(len(assignments["fast"]), len(assignments["slow"]))
-        self.assertEqual(sorted(assignments["fast"] + assignments["slow"]), list(range(12)))
-
-    async def test_failure_timeout_and_cancel_clean_only_current_assignments(self):
-        async def dispatch(worker, index):
-            return {"index": index, "prompt_id": f"{worker}-{index}"}
-        cancelled = []
-        async def cleanup(worker, handle):
-            cancelled.append((worker, handle["prompt_id"]))
-        async def timeout(_inflight):
-            raise TimeoutError("worker timed out")
-        with self.assertRaisesRegex(TimeoutError, "timed out"):
-            await dynamic_schedule(4, ["one", "two"], dispatch, timeout, cleanup)
-        self.assertCountEqual(cancelled, [("one", "one-0"), ("two", "two-1")])
-
-        cancelled.clear()
-        async def wrong_index(inflight):
-            worker, handle = next(iter(inflight.items()))
-            return worker, handle["index"] + 99, "bad"
-        with self.assertRaisesRegex(RuntimeError, "unexpected image index"):
-            await dynamic_schedule(2, ["one"], dispatch, wrong_index, cleanup)
-        self.assertEqual(len(cancelled), 1)
-
-        cancelled.clear()
-        with self.assertRaises(ImageListCancelledError):
-            await dynamic_schedule(2, ["one", "two"], dispatch, timeout, cleanup,
-                                   interrupted=lambda: True)
-        self.assertEqual(len(cancelled), 0)
-
-    async def test_no_workers_and_duplicate_results_fail_closed(self):
-        async def unused(*_args):
-            self.fail("must not dispatch without workers")
-        with self.assertRaisesRegex(RuntimeError, "No free"):
-            await dynamic_schedule(2, [], unused, unused, unused)
-        async def dispatch(worker, index):
-            return {"index": index}
-        async def duplicate(inflight):
-            worker, _handle = next(iter(inflight.items()))
-            return worker, 0, "same"
-        with self.assertRaisesRegex(RuntimeError, "unexpected image index"):
-            await dynamic_schedule(2, ["one"], dispatch, duplicate, lambda *_: asyncio.sleep(0))
-
-
-if __name__ == "__main__":
-    unittest.main()

tests/test_image_list_integration_gaps.py

@@ -1,180 +0,0 @@
-"""Regression coverage at the scheduler and real HTTP transport boundaries."""
-import asyncio
-import importlib.util
-import json
-from pathlib import Path
-import sys
-import types
-import unittest
-from unittest.mock import patch
-
-import torch
-from aiohttp import web
-from aiohttp.test_utils import TestClient, TestServer
-
-from image_list import runtime
-from test_image_list import graph_two_sequential_regions, start_node
-
-
-def load_routes():
-    root = Path(__file__).resolve().parents[1]
-    name = '_image_list_transport_tests'
-    routes = web.RouteTableDef()
-    modules = {
-        name: types.ModuleType(name),
-        name + '.api': types.ModuleType(name + '.api'),
-        name + '.image_list': types.ModuleType(name + '.image_list'),
-        name + '.image_list.runtime': runtime,
-        'server': types.SimpleNamespace(PromptServer=types.SimpleNamespace(
-            instance=types.SimpleNamespace(routes=routes))),
-    }
-    spec = importlib.util.spec_from_file_location(name + '.api.image_list_routes', root / 'api/image_list_routes.py')
-    module = importlib.util.module_from_spec(spec)
-    with patch.dict(sys.modules, modules):
-        spec.loader.exec_module(module)
-    return module, routes
-
-
-class TransportTests(unittest.IsolatedAsyncioTestCase):
-    async def asyncSetUp(self):
-        self.module, routes = load_routes()
-        app = web.Application()
-        app.add_routes(routes)
-        self.client = TestClient(TestServer(app))
-        await self.client.start_server()
-
-    async def asyncTearDown(self):
-        await self.client.close()
-        runtime._WORKER_ITEMS.clear()
-
-    async def test_fragmented_chunked_item_upload_is_consumable(self):
-        token = runtime.create_worker_task_token()
-        item = {'image': torch.full((1, 2, 3, 3), 0.25), 'index': 2}
-        data = json.dumps({'token': token, 'context_id': 'context', 'item': runtime.encode_item(item)}).encode()
-        async def chunks():
-            yield data[:8]
-            await asyncio.sleep(0.02)
-            yield data[8:]
-        response = await self.client.post('/distributed/image-list/item', data=chunks())
-        self.assertEqual(response.status, 200, await response.text())
-        result = runtime.consume_worker_item(token)
-        self.assertTrue(torch.equal(result['image'], item['image']))
-        self.assertEqual(result['index'], 2)
-
-    async def test_complete_item_upload_and_replay_rejection(self):
-        token = runtime.create_worker_task_token()
-        data = {'token': token, 'context_id': 'context', 'item': runtime.encode_item(
-            {'image': torch.zeros((1, 2, 3, 3)), 'index': 0})}
-        response = await self.client.post('/distributed/image-list/item', json=data)
-        self.assertEqual(response.status, 200, await response.text())
-        response = await self.client.post('/distributed/image-list/item', json=data)
-        self.assertEqual(response.status, 400)
-        runtime.consume_worker_item(token)
-        with self.assertRaisesRegex(RuntimeError, 'already consumed'):
-            runtime.consume_worker_item(token)
-
-    async def test_chunked_body_limit_is_aggregate(self):
-        async def chunks():
-            yield b'{"context_id":"'
-            await asyncio.sleep(0.01)
-            yield b'x' * 60
-            yield b'"}'
-        with patch.object(self.module, 'MAX_REQUEST_BYTES', 64):
-            response = await self.client.post('/distributed/image-list/item', data=chunks())
-            self.assertEqual(response.status, 400)
-            self.assertIn('exceeds', await response.text())
-
-    async def test_invalid_result_shapes_are_rejected(self):
-        for kind, shape in [('image', (1, 2, 3)), ('mask', (1, 2, 3, 1))]:
-            result = runtime.encode_result({'image': torch.zeros((1, 2, 3, 3)),
-                                             'mask': torch.zeros((1, 2, 3))})
-            result[kind]['shape'] = list(shape)
-            with self.assertRaises(ValueError):
-                runtime.decode_result(result)
-
-
-class SchedulingBoundaryTests(unittest.IsolatedAsyncioTestCase):
-    async def test_user_interrupt_during_result_wait_cleans_up_promptly(self):
-        entered = asyncio.Event()
-        cancelled = asyncio.Event()
-        wait_cancelled = asyncio.Event()
-        dispatches, cleanups = [], []
-        async def dispatch(worker, index):
-            dispatches.append(index)
-            return {'index': index}
-        async def wait_next(_inflight):
-            entered.set()
-            try:
-                await asyncio.Event().wait()
-            finally:
-                wait_cancelled.set()
-        async def cleanup(worker, handle):
-            cleanups.append((worker, handle['index']))
-        task = asyncio.create_task(runtime.dynamic_schedule(
-            3, ['one'], dispatch, wait_next, cleanup, cancelled.is_set))
-        await entered.wait()
-        cancelled.set()
-        with self.assertRaises(runtime.ImageListCancelledError):
-            await asyncio.wait_for(task, timeout=1)
-        self.assertTrue(wait_cancelled.is_set())
-        self.assertEqual(dispatches, [0])
-        self.assertEqual(cleanups, [('one', 0)])
-
-    async def test_already_interrupted_does_not_submit_any_work(self):
-        async def unused(*args):
-            self.fail('pre-cancelled job must not dispatch or wait')
-        with self.assertRaises(runtime.ImageListCancelledError):
-            await runtime.dynamic_schedule(2, ['one'], unused, unused, unused, lambda: True)
-
-    async def test_task_cancel_during_submit_waits_for_owned_handle(self):
-        entered, release = asyncio.Event(), asyncio.Event()
-        cleanups = []
-        async def dispatch(worker, index):
-            entered.set()
-            await release.wait()
-            return {'index': index, 'prompt_id': 'owned'}
-        async def unused(*args):
-            self.fail('cancelled dispatch must not enter result wait')
-        async def cleanup(worker, handle):
-            cleanups.append(handle['prompt_id'])
-        task = asyncio.create_task(runtime.dynamic_schedule(2, ['one'], dispatch, unused, cleanup))
-        await entered.wait()
-        task.cancel()
-        release.set()
-        with self.assertRaises(asyncio.CancelledError):
-            await task
-        self.assertEqual(cleanups, ['owned'])
-
-    def test_output_nodes_inside_region_are_rejected(self):
-        prompt = graph_two_sequential_regions()
-        output_class = type('OutputOperation', (), {'OUTPUT_NODE': True})
-        registry = types.SimpleNamespace(NODE_CLASS_MAPPINGS={'FirstFrameOperation': output_class})
-        with patch.dict(sys.modules, {'nodes': registry}):
-            with self.assertRaisesRegex(ValueError, 'must be placed after'):
-                runtime.validate_image_list_regions(prompt)
-
-    def test_shared_scalars_are_validated_before_dispatch(self):
-        for value in (float('nan'), object(), torch.zeros(2, dtype=torch.int32)):
-            with self.assertRaises(TypeError):
-                runtime.normalize_image_list_inputs(torch.zeros((1, 2, 2, 3)), values={'value_1': [value]})
-
-    def test_three_sequential_regions_can_share_worker_loader(self):
-        prompt = graph_two_sequential_regions()
-        prompt.pop('8')
-        prompt.update({
-            '8': start_node(['7', 0]),
-            '9': {'class_type': 'ThirdOperation', 'inputs': {'image': ['8', 0]}},
-            '10': {'class_type': runtime.END_CLASS,
-                   'inputs': {'images': ['9', 0], 'context': ['8', runtime.CONTEXT_OUT]}},
-            '11': {'class_type': 'GenericModelLoader', 'inputs': {}},
-        })
-        for node_id in ['3', '6', '9']:
-            prompt[node_id]['inputs']['model'] = ['11', 0]
-        # Shared dependencies are verified against the registered output schema.
-        loader = type('GenericModelLoader', (), {'RETURN_TYPES': ('MODEL',)})
-        registry = types.SimpleNamespace(NODE_CLASS_MAPPINGS={'GenericModelLoader': loader})
-        with patch.dict(sys.modules, {'nodes': registry}):
-            regions = runtime.validate_image_list_regions(prompt)
-        self.assertEqual(len(regions), 3)
-        for region in regions:
-            self.assertIn('11', region['node_ids'])

tests/test_image_list_lifecycle.py

@@ -1,177 +0,0 @@
-"""Tests for periodic image-list server lifecycle cleanup."""
-import asyncio
-import time
-import unittest
-from types import SimpleNamespace
-from typing import Any, cast
-from unittest.mock import patch
-
-from aiohttp import web
-
-from image_list import lifecycle, runtime
-
-
-class ImageListLifecycleTests(unittest.IsolatedAsyncioTestCase):
-    def setUp(self):
-        with runtime._LOCK:
-            runtime._MASTER_JOBS.clear()
-            runtime._WORKER_ITEMS.clear()
-
-    def tearDown(self):
-        with runtime._LOCK:
-            runtime._MASTER_JOBS.clear()
-            runtime._WORKER_ITEMS.clear()
-
-    @staticmethod
-    def _state(prompt_id=None, *, active=False, age=10_000.0):
-        return SimpleNamespace(
-            prompt_id=prompt_id,
-            active=active,
-            created_at=time.monotonic() - age,
-        )
-
-    @staticmethod
-    def _server(queue):
-        return SimpleNamespace(app=web.Application(), prompt_queue=queue)
-
-    async def _start(self, server, *, interval=0.01, grace=0.01):
-        lifecycle.install_lifecycle(
-            server,
-            runtime,
-            interval_seconds=interval,
-            orphan_grace_seconds=grace,
-        )
-        runner = web.AppRunner(server.app)
-        await runner.setup()
-        return runner
-
-    async def test_unattended_sweep_removes_orphan_and_expired_staging(self):
-        context_id = "orphan"
-        with runtime._LOCK:
-            cast(Any, runtime._MASTER_JOBS)[context_id] = self._state()
-            runtime._WORKER_ITEMS["expired-token"] = (time.monotonic() - 1, {"image": "stale"})
-
-        queue = SimpleNamespace(
-            get_current_queue_volatile=lambda: ([], []),
-        )
-        server = self._server(queue)
-        swept = asyncio.Event()
-        sweep_worker_items = runtime._sweep_worker_items
-
-        def track_worker_sweep():
-            swept.set()
-            sweep_worker_items()
-
-        with patch.object(runtime, "_sweep_worker_items", side_effect=track_worker_sweep):
-            runner = await self._start(server)
-            try:
-                await asyncio.wait_for(swept.wait(), timeout=0.5)
-                self.assertNotIn(context_id, runtime._MASTER_JOBS)
-                self.assertNotIn("expired-token", runtime._WORKER_ITEMS)
-                with self.assertRaisesRegex(ValueError, "Unknown or expired"):
-                    runtime.set_master_result(context_id, 0, "late-token", {})
-            finally:
-                await runner.cleanup()
-
-    async def test_old_active_and_queued_contexts_survive_but_orphans_expire(self):
-        states = {
-            "active": self._state("active-prompt", active=True),
-            "queued": self._state("queued-prompt"),
-            "orphan": self._state("missing-prompt"),
-            "unbound": self._state(None),
-        }
-        with runtime._LOCK:
-            cast(Any, runtime._MASTER_JOBS).update(states)
-        queue = SimpleNamespace(
-            get_current_queue_volatile=lambda: (
-                [[1, "active-prompt", {}]],
-                [(2, "queued-prompt", {}, {}, None, None)],
-            ),
-        )
-        runner = await self._start(self._server(queue))
-        try:
-            await asyncio.sleep(0.04)
-            self.assertIs(runtime._MASTER_JOBS.get("active"), states["active"])
-            self.assertIs(runtime._MASTER_JOBS.get("queued"), states["queued"])
-            self.assertNotIn("orphan", runtime._MASTER_JOBS)
-            self.assertNotIn("unbound", runtime._MASTER_JOBS)
-        finally:
-            await runner.cleanup()
-
-    async def test_falls_back_to_nonvolatile_queue_reader(self):
-        state = self._state("queued-prompt")
-        with runtime._LOCK:
-            cast(Any, runtime._MASTER_JOBS)["queued"] = state
-        fallback_reads = 0
-
-        def volatile_failure():
-            raise RuntimeError("volatile view unavailable")
-
-        def current_queue():
-            nonlocal fallback_reads
-            fallback_reads += 1
-            return ([], [[1, "queued-prompt", {}]])
-
-        queue = SimpleNamespace(
-            get_current_queue_volatile=volatile_failure,
-            get_current_queue=current_queue,
-        )
-        runner = await self._start(self._server(queue))
-        try:
-            await asyncio.sleep(0.04)
-            self.assertGreaterEqual(fallback_reads, 1)
-            self.assertIs(runtime._MASTER_JOBS.get("queued"), state)
-        finally:
-            await runner.cleanup()
-
-    async def test_queue_read_error_protects_bound_contexts(self):
-        with runtime._LOCK:
-            cast(Any, runtime._MASTER_JOBS)["bound"] = self._state("bound-prompt")
-            cast(Any, runtime._MASTER_JOBS)["unbound"] = self._state(None)
-
-        def fail_read():
-            raise RuntimeError("queue unavailable")
-
-        queue = SimpleNamespace(
-            get_current_queue_volatile=fail_read,
-            get_current_queue=fail_read,
-        )
-        runner = await self._start(self._server(queue))
-        try:
-            await asyncio.sleep(0.04)
-            self.assertIn("bound", runtime._MASTER_JOBS)
-            self.assertNotIn("unbound", runtime._MASTER_JOBS)
-        finally:
-            await runner.cleanup()
-
-    async def test_install_and_removal_are_idempotent_and_task_stops_on_cleanup(self):
-        queue_reads = 0
-
-        def read_queue():
-            nonlocal queue_reads
-            queue_reads += 1
-            return ([], [])
-
-        server = self._server(SimpleNamespace(get_current_queue_volatile=read_queue))
-        lifecycle.install_lifecycle(server, runtime, interval_seconds=0.01, orphan_grace_seconds=0.01)
-        lifecycle.install_lifecycle(server, runtime, interval_seconds=0.01, orphan_grace_seconds=0.01)
-        self.assertEqual(len(server.app.cleanup_ctx), 1)
-
-        with runtime._LOCK:
-            cast(Any, runtime._MASTER_JOBS)["removed"] = self._state()
-        runtime.remove_master_job("removed")
-        runtime.remove_master_job("removed")
-        self.assertNotIn("removed", runtime._MASTER_JOBS)
-
-        runner = web.AppRunner(server.app)
-        await runner.setup()
-        await asyncio.sleep(0.035)
-        self.assertGreaterEqual(queue_reads, 1)
-        await runner.cleanup()
-        reads_after_cleanup = queue_reads
-        await asyncio.sleep(0.035)
-        self.assertEqual(queue_reads, reads_after_cleanup)
-
-
-if __name__ == "__main__":
-    unittest.main()

tests/test_image_list_master_runtime.py

@@ -1,437 +0,0 @@
-"""Master slot regressions, plus an opt-in real CPU ComfyUI subprocess probe."""
-import asyncio
-import importlib
-import json
-import os
-from pathlib import Path
-import subprocess
-import sys
-import types
-import unittest
-from unittest.mock import AsyncMock, patch
-
-import torch
-
-ROOT = Path(__file__).resolve().parents[1]
-
-
-def graph():
-    return {
-        'source': {'class_type': 'Source', 'inputs': {}},
-        'start': {'class_type': 'DistributedImageListStart', 'inputs': {'images': ['source', 0]}},
-        'op': {'class_type': 'Operation', 'inputs': {'image': ['start', 0]}},
-        'end': {'class_type': 'DistributedImageListEnd', 'inputs': {
-            'images': ['op', 0], 'context': ['start', 3]}},
-    }
-
-
-class Response:
-    status = 200
-
-    def __init__(self, data=None):
-        self.data = data
-
-    async def __aenter__(self):
-        return self
-
-    async def __aexit__(self, *args):
-        return None
-
-    async def json(self):
-        return self.data
-
-    def raise_for_status(self):
-        pass
-
-
-class MasterRuntimeTests(unittest.IsolatedAsyncioTestCase):
-    def setUp(self):
-        name = '_master_slot_tests'
-        self.modules = patch.dict(sys.modules)
-        self.modules.start()
-        self.addCleanup(self.modules.stop)
-        for suffix in ('', '.image_list', '.api', '.api.orchestration', '.utils'):
-            module = types.ModuleType(name + suffix)
-            module.__path__ = [str(ROOT / suffix.lstrip('.').replace('.', '/'))]
-            sys.modules[name + suffix] = module
-        self.rt = importlib.import_module(name + '.image_list.runtime')
-        self.workers = []
-        self.preflight = AsyncMock(side_effect=lambda workers, **kw: {w['id']: True for w in workers})
-        self.dispatched = []
-        self.local_calls = []
-        self.tokens = []
-        self.closed = 0
-        self.local_error = None
-        self.local_delay = 0.005
-        self.interrupted = False
-        self.local_started = asyncio.Event()
-        self.local_finished = asyncio.Event()
-        self.deadline = 0
-        self.slot_instances = []
-        self.config = {'workers': self.workers, 'settings': {}}
-        fixture = self
-
-        class FakeSlot:
-            def __init__(self, model_callback=None):
-                fixture.slot_instances.append(self)
-                self.tasks = []
-
-            def submit(self, prompt, prompt_id, end_id, workflow_meta=None):
-                async def execute():
-                    inputs = prompt[end_id]['inputs']
-                    index = inputs['worker_item_index']
-                    token = inputs['worker_result_token']
-                    fixture.local_calls.append(index)
-                    fixture.tokens.append(token)
-                    fixture.local_started.set()
-                    try:
-                        await asyncio.sleep(fixture.local_delay)
-                        if fixture.local_error:
-                            raise fixture.local_error
-                        item = fixture.rt.consume_worker_item(token)
-                        fixture.rt.set_master_result(fixture.context, index, token, {
-                            'image': item['image'], 'mask': item.get('mask'),
-                            'text': item['values'].get('value_1')})
-                    finally:
-                        fixture.local_finished.set()
-                task = asyncio.create_task(execute())
-                self.tasks.append(task)
-                return task
-
-            def request_stop(self):
-                pass
-
-            async def close(self):
-                # Model a synchronous node that must finish before its thread can join.
-                await asyncio.gather(*self.tasks, return_exceptions=True)
-                fixture.closed += 1
-
-        dependencies = {
-            '.image_list.local_executor': types.SimpleNamespace(LocalExecutionSlot=FakeSlot),
-            '.api.orchestration.dispatch': types.SimpleNamespace(
-                preflight_image_list_workers=self.preflight,
-                dispatch_worker_prompt=AsyncMock(side_effect=self.dispatch)),
-            '.api.orchestration.prompt_transform': types.SimpleNamespace(
-                _resolve_direct_workers_for_image_list=lambda config, ids: self.workers),
-            '.api.orchestration.media_sync': types.SimpleNamespace(
-                sync_worker_media=AsyncMock(), fetch_worker_path_separator=AsyncMock(return_value=None)),
-            '.utils.config': types.SimpleNamespace(load_config=lambda: self.config,
-                get_worker_timeout_seconds=lambda: 60,
-                get_image_list_job_timeout_seconds=lambda: self.deadline),
-            '.utils.network': types.SimpleNamespace(
-                build_master_callback_url=lambda *a, **kw: 'http://master',
-                build_worker_url=lambda worker, path: 'http://remote' + path,
-                get_client_session=AsyncMock(return_value=self)),
-            '.utils.logging': types.SimpleNamespace(log=lambda *args: None),
-        }
-        for suffix, module in dependencies.items():
-            sys.modules[name + suffix] = module
-        self.network = dependencies['.utils.network']
-        sys.modules['aiohttp'] = types.SimpleNamespace(ClientTimeout=lambda **kw: kw, ClientError=ConnectionError)
-        comfy = types.ModuleType('comfy')
-        comfy.model_management = types.SimpleNamespace(processing_interrupted=lambda: self.interrupted)
-        sys.modules['comfy'] = comfy
-        sys.modules['comfy.model_management'] = comfy.model_management
-
-    def create(self, participates=True, count=5):
-        self.items = [{'image': torch.full((1, i + 1, 2, 3), float(i)),
-                       'mask': torch.full((1, i + 1, 2), float(i)), 'values': {'value_1': str(i)},
-                       'index': i}
-                      for i in range(count)]
-        self.context = self.rt.create_master_job('start', 'end', graph(), self.items,
-            [w['id'] for w in self.workers], '', master_participates=participates)
-        self.state = self.rt.get_master_job(self.context)
-        return self.context
-
-    def post(self, url, *, data=None, **kwargs):
-        if url.endswith('/distributed/image-list/item'):
-            payload = json.loads(data)
-            self.rt.store_worker_item(payload['token'], payload['item'])
-        return Response({})
-
-    def delete(self, url, **kwargs):
-        self.rt.delete_worker_item(url.rsplit('/', 1)[-1])
-        return Response({})
-
-    def get(self, url, **kwargs):
-        pid = url.rsplit('/', 1)[-1]
-        return Response({pid: {'status': {'completed': True, 'status_str': 'success'}}})
-
-    async def dispatch(self, worker, prompt, meta, **kwargs):
-        inputs = prompt['end']['inputs']
-        index, token = inputs['worker_item_index'], inputs['worker_result_token']
-        self.dispatched.append(index)
-        # A remote callback is ready immediately; the local slot is slower.
-        item = self.rt.consume_worker_item(token)
-        self.rt.set_master_result(self.context, index, token, {
-            'image': item['image'], 'mask': item['mask'], 'text': item['values']['value_1']})
-        return f'remote-{index}'
-
-    async def test_local_only_preserves_order_and_uses_no_http(self):
-        result = await self.rt.run_image_list_job(self.create())
-        self.assertEqual(result[2], [str(i) for i in range(5)])
-        self.assertEqual(self.local_calls, list(range(5)))
-        self.assertEqual(len(self.slot_instances), 1)
-        self.assertEqual(len(set(self.tokens)), 5)
-        self.network.get_client_session.assert_not_awaited()
-        self.assertEqual(self.closed, 1)
-        self.assertFalse(self.rt._WORKER_ITEMS)
-        self.assertFalse(self.rt._MASTER_JOBS)
-        for i, (image, mask) in enumerate(zip(*result[:2])):
-            self.assertTrue(torch.equal(image, self.items[i]['image']))
-            self.assertTrue(torch.equal(mask, self.items[i]['mask']))
-
-    async def test_mixed_pool_claims_from_one_queue_and_relaxes_remote_idle(self):
-        self.workers.append({'id': 'master', 'type': 'local'})  # Not the in-process master slot.
-        self.local_delay = 0.04
-        result = await self.rt.run_image_list_job(self.create(count=9))
-        self.assertTrue(self.local_calls)
-        self.assertTrue(self.dispatched)
-        self.assertEqual(sorted(self.local_calls + self.dispatched), list(range(9)))
-        self.assertEqual(result[2], [str(i) for i in range(9)])
-        self.preflight.assert_awaited_once_with(self.workers, require_idle=False)
-
-    async def test_orchestrator_mode_never_instantiates_local_slot(self):
-        self.workers.append({'id': 'worker', 'type': 'remote'})
-        await self.rt.run_image_list_job(self.create(participates=False))
-        self.assertFalse(self.slot_instances)
-        self.preflight.assert_awaited_once_with(self.workers, require_idle=True)
-        with self.assertRaisesRegex(RuntimeError, 'worker|fallback'):
-            self.rt.create_master_job('start', 'end', graph(), self.items, [], '')
-
-    async def test_local_failure_surfaces_and_reclaims_staged_token(self):
-        self.local_error = RuntimeError('local branch failed')
-        with self.assertRaisesRegex(RuntimeError, 'local branch failed'):
-            await self.rt.run_image_list_job(self.create())
-        self.assertEqual(self.closed, 1)
-        self.assertFalse(self.rt._WORKER_ITEMS)
-        self.assertFalse(self.rt._MASTER_JOBS)
-
-    async def test_deadline_joins_local_execution_before_removing_context(self):
-        self.deadline, self.local_delay = 0.03, 0.1
-        with self.assertRaisesRegex(TimeoutError, 'total deadline'):
-            await self.rt.run_image_list_job(self.create())
-        self.assertTrue(self.local_finished.is_set())
-        self.assertEqual(self.closed, 1)
-        self.assertFalse(self.rt._WORKER_ITEMS)
-        self.assertFalse(self.rt._MASTER_JOBS)
-
-
-    async def test_cancelled_scheduler_joins_owned_execution_even_if_cancelled_again(self):
-        self.local_delay = 0.12
-        task = asyncio.create_task(self.rt.run_image_list_job(self.create()))
-        await self.local_started.wait()
-        task.cancel()
-        await asyncio.sleep(0.02)
-        task.cancel()
-        with self.assertRaises(self.rt.ImageListCancelledError):
-            await task
-        self.assertTrue(self.local_finished.is_set())
-        self.assertFalse(self.rt._WORKER_ITEMS)
-        self.assertFalse(self.rt._MASTER_JOBS)
-
-
-class RealFrameworkTests(unittest.TestCase):
-    @unittest.skipUnless(os.environ.get('COMFYUI_SOURCE_ROOT'), 'set COMFYUI_SOURCE_ROOT for the real CPU framework probe')
-    def test_real_framework_local_slot(self):
-        result = subprocess.run([sys.executable, str(Path(__file__).resolve()), '--real-probe',
-                                 os.environ['COMFYUI_SOURCE_ROOT']], text=True, capture_output=True, timeout=60)
-        self.assertEqual(result.returncode, 0, result.stdout + result.stderr)
-        self.assertIn('MASTER_SLOT_REAL_OK', result.stdout)
-
-
-def real_framework_probe(comfy_root):
-    import threading
-    import time
-    sys.path.insert(0, str(Path(comfy_root).resolve()))
-    import comfy.cli_args
-    comfy.cli_args.args.cpu = True
-    comfy.cli_args.args.disable_assets = True
-    import execution
-    import nodes
-    import comfy.memory_management as memory
-    import comfy.model_management as models
-    import comfy.utils as utils
-    from comfy_execution import progress
-    from comfy_execution.graph_utils import GraphBuilder
-    from comfy_api.v0_0_2 import io
-
-    name = '_real_master_slot'
-    for suffix in ('', '.image_list', '.utils', '.nodes'):
-        module = types.ModuleType(name + suffix)
-        module.__path__ = [str(ROOT / suffix.lstrip('.').replace('.', '/'))]
-        sys.modules[name + suffix] = module
-    # No HTTP/service/extension bootstrap is needed by the local slot. Fail if
-    # either boundary tries to use a server-loop bridge instead of direct memory.
-    def forbidden(*args, **kwargs):
-        raise AssertionError('local boundary used the server loop or /prompt')
-    sys.modules[name + '.utils.async_helpers'] = types.SimpleNamespace(run_async_in_server_loop=forbidden)
-    config = {'settings': {}}
-    sys.modules[name + '.utils.config'] = types.SimpleNamespace(load_config=lambda: config)
-    rt = importlib.import_module(name + '.image_list.runtime')
-    local = importlib.import_module(name + '.image_list.local_executor')
-    boundary = importlib.import_module(name + '.nodes.image_list')
-    nodes.NODE_CLASS_MAPPINGS.update(DistributedImageListStart=boundary.DistributedImageListStart,
-                                   DistributedImageListEnd=boundary.DistributedImageListEnd)
-    loads, calls = [], []
-    entered, stopped = threading.Event(), threading.Event()
-    mode = ['normal']
-    model = object()
-
-    class Loader:
-        @classmethod
-        def INPUT_TYPES(cls):
-            return {'required': {}}
-        RETURN_TYPES = ('MODEL',)
-        FUNCTION = 'run'
-        def run(self):
-            loads.append(True)
-            return (model,)
-
-    class Operation:
-        @classmethod
-        def INPUT_TYPES(cls):
-            return {'required': {'image': ('IMAGE',), 'model': ('MODEL',),
-                                 'mask': ('MASK',), 'text': ('STRING', {'forceInput': True})}}
-        RETURN_TYPES = ('IMAGE', 'MASK', 'STRING')
-        FUNCTION = 'run'
-        async def run(self, image, model, mask, text):
-            calls.append(text)
-            entered.set()
-            if mode[0] == 'error':
-                raise ValueError('probe branch failure')
-            if mode[0] in ('cancel', 'interrupt'):
-                try:
-                    while True:
-                        utils.ProgressBar(1).update(1)
-                        await asyncio.sleep(0.01)
-                finally:
-                    await asyncio.sleep(0.02)
-                    stopped.set()
-            if mode[0] == 'sync':
-                # No progress/interrupt checks: cleanup still must join, not
-                # release the parent just because an asyncio Future is cancelled.
-                time.sleep(0.08)
-                stopped.set()
-            await asyncio.sleep(0)
-            return image + 1, mask + 1, text + ':done'
-
-    class Expanded:
-        @classmethod
-        def INPUT_TYPES(cls):
-            return Operation.INPUT_TYPES()
-        RETURN_TYPES = Operation.RETURN_TYPES
-        FUNCTION = 'run'
-        def run(self, image, model, mask, text):
-            builder = GraphBuilder()
-            node = builder.node('Operation', image=image, model=model, mask=mask, text=text)
-            return {'result': (node.out(0), node.out(1), node.out(2)), 'expand': builder.finalize()}
-
-    nodes.NODE_CLASS_MAPPINGS.update(Loader=Loader, Operation=Operation, Expanded=Expanded)
-    prompt = graph()
-    prompt['loader'] = {'class_type': 'Loader', 'inputs': {}}
-    prompt['op']['inputs'].update(model=['loader', 0], mask=['start', rt.MASK_OUT],
-                                  text=['start', rt.VALUE_OUT_BASE])
-    prompt['end']['inputs'].update(masks=['op', 1], texts=['op', 2])
-    items = [{'image': torch.full((1, i + 1, 2, 3), float(i)),
-              'mask': torch.full((1, i + 1, 2), float(i)), 'values': {'value_1': str(i)}, 'index': i}
-             for i in range(3)]
-    region = rt.validate_image_list_regions(prompt)[0]
-    parent_progress = progress.ProgressRegistry('outer', execution.DynamicPrompt({}))
-    parent_progress.nodes['outer'] = {'state': progress.NodeState.Running, 'value': 3, 'max': 7}
-    progress.global_progress_registry = parent_progress
-    parent_hook = lambda *a, **kw: None
-    utils.PROGRESS_BAR_HOOK = parent_hook
-    ram_callback = lambda *a, **kw: None
-    memory.set_ram_cache_release_state(ram_callback, 123)
-    GraphBuilder.set_default_prefix('outer-prefix', 7, 9)
-
-    def restored():
-        assert progress.global_progress_registry is parent_progress
-        assert parent_progress.nodes == {'outer': {'state': progress.NodeState.Running, 'value': 3, 'max': 7}}
-        assert utils.PROGRESS_BAR_HOOK is parent_hook
-        assert memory.extra_ram_release_callback is ram_callback and memory.RAM_CACHE_HEADROOM == 123
-        assert (GraphBuilder._default_prefix_root, GraphBuilder._default_prefix_call_index,
-                GraphBuilder._default_prefix_graph_index) == ('outer-prefix', 7, 9)
-        assert not any(t.name.startswith('image-list-master') for t in threading.enumerate())
-
-    async def probe():
-        # Verify actual V3 input parsing and list mapping, not only direct execute.
-        for supplied, delegated, expected in [(None, False, True), (None, True, False),
-                                               (False, False, False), (True, True, True)]:
-            config['settings']['master_delegate_only'] = delegated
-            inputs = {'image_list_enabled_worker_ids': '["remote"]'}
-            if supplied is not None:
-                inputs['image_list_master_participates'] = supplied
-            cls = boundary.DistributedImageListStart
-            data, _, v3 = execution.get_input_data(inputs, cls, 'start', dynprompt=execution.DynamicPrompt(prompt))
-            data['images'] = [torch.zeros((3, 2, 2, 3))]
-            result = await execution._async_map_node_over_list('outer', 'start', cls, data, cls.FUNCTION, v3_data=v3)
-            result = await execution.resolve_map_node_over_list_results(result)
-            context = result[0].result[rt.CONTEXT_OUT][0]
-            assert rt.get_master_job(context).master_participates is expected
-            rt.remove_master_job(context)
-
-        for behavior in ('normal', 'expand', 'error', 'cancel', 'interrupt', 'sync', 'already-interrupted'):
-            mode[0] = behavior
-            prompt['op']['class_type'] = 'Expanded' if behavior == 'expand' else 'Operation'
-            loads.clear()
-            calls.clear()
-            entered.clear()
-            stopped.clear()
-            tracked = []
-            slot = local.LocalExecutionSlot(model_callback=tracked.append)
-            context = rt.create_master_job('start', 'end', prompt, items, [], '', master_participates=True)
-            try:
-                for index in range(3 if behavior in ('normal', 'expand') else 1):
-                    token = rt.create_worker_task_token()
-                    rt.register_assignment(context, index, 'local', token)
-                    rt._WORKER_ITEMS[token] = (time.monotonic() + 60, items[index])
-                    worker_prompt = rt.build_worker_prompt(prompt, region, token, context, index, '')
-                    for key in ('start', 'end'):
-                        worker_prompt[key]['inputs']['image_list_local_execution'] = True
-                    if behavior == 'already-interrupted':
-                        models.interrupt_current_processing(True)
-                    future = slot.submit(worker_prompt, f'local-{index}', 'end')
-                    if behavior in ('cancel', 'interrupt', 'sync'):
-                        assert await asyncio.to_thread(entered.wait, 5)
-                        if behavior == 'interrupt':
-                            models.interrupt_current_processing(True)
-                        else:
-                            slot.request_stop()
-                    try:
-                        await asyncio.wait_for(asyncio.shield(asyncio.wrap_future(future)), 5)
-                    except rt.ImageListCancelledError:
-                        assert behavior in ('cancel', 'interrupt', 'sync', 'already-interrupted'), behavior
-                    except RuntimeError as exc:
-                        assert behavior == 'error' and 'probe branch failure' in str(exc), str(exc)
-                    else:
-                        assert behavior in ('normal', 'expand'), behavior
-                        result = rt.get_master_job(context).results[index]
-                        assert torch.equal(result['image'], items[index]['image'] + 1)
-                        assert torch.equal(result['mask'], items[index]['mask'] + 1)
-                        assert result['text'] == f'{index}:done'
-                if behavior in ('normal', 'expand'):
-                    assert len(loads) == 1 and calls == ['0', '1', '2'], (loads, calls)
-                    assert any(model in outputs for outputs in tracked), tracked
-                if behavior in ('cancel', 'interrupt', 'sync'):
-                    assert stopped.is_set(), behavior
-                if behavior == 'already-interrupted':
-                    assert models.processing_interrupted() and not calls
-            finally:
-                await slot.close()
-                for token in rt.get_master_job(context).expected_tokens.values():
-                    rt.delete_worker_item(token)
-                rt.remove_master_job(context)
-                models.interrupt_current_processing(False)  # Only test setup/teardown.
-            restored()
-        assert not rt._MASTER_JOBS and not rt._WORKER_ITEMS
-    asyncio.run(probe())
-    print('MASTER_SLOT_REAL_OK V3 metadata; cache reuse; async/expanded graph; errors; cancellation; state restoration')
-
-
-if __name__ == '__main__':
-    if len(sys.argv) > 1 and sys.argv[1] == '--real-probe':
-        real_framework_probe(sys.argv[2])
-    else:
-        unittest.main()

tests/test_image_list_master_selection.py

@@ -1,114 +0,0 @@
-"""Image-list participation uses the existing master orchestrator setting."""
-import json
-import sys
-import unittest
-from unittest.mock import AsyncMock, patch
-
-from test_serverless_integration import PKG, load_modules
-
-
-class MasterSelectionTests(unittest.IsolatedAsyncioTestCase):
-    async def asyncSetUp(self):
-        self.q, modules, self.ps = load_modules()
-        self.modules_patch = patch.dict(sys.modules, modules)
-        self.modules_patch.start()
-        self.addCleanup(self.modules_patch.stop)
-        self.worker = {"id": "worker", "name": "Worker", "type": "local", "enabled": True}
-        self.config = {"workers": [self.worker], "settings": {}}
-        self.q.load_config = lambda: self.config
-        # Model the legacy helper's fallback: image-list must not inherit it.
-        self.q.select_active_workers = AsyncMock(side_effect=lambda workers, ws, delegate, **kw:
-                                                (workers, delegate if workers else False))
-        self.preflight = AsyncMock(return_value={"worker": True})
-        sys.modules[PKG + ".api.orchestration.dispatch"].preflight_image_list_workers = self.preflight
-        self.q.queue_prompt_payload = AsyncMock(return_value={"prompt_id": "master-prompt", "number": 1})
-        self.prompt = {
-            "1": {"class_type": "EmptyImage", "inputs": {}},
-            "2": {"class_type": "DistributedImageListStart", "inputs": {"images": ["1", 0]}},
-            "3": {"class_type": "DistributedImageListEnd", "inputs": {"images": ["2", 0], "context": ["2", 3]}},
-        }
-
-    async def queue(self, **kwargs):
-        return await self.q.orchestrate_distributed_execution(self.prompt, None, "client", **kwargs)
-
-    def prepared_inputs(self):
-        return self.q.queue_prompt_payload.call_args.args[0]["2"]["inputs"]
-
-    async def test_default_includes_master_and_direct_workers(self):
-        result = await self.queue()
-        self.assertIs(self.prepared_inputs()["image_list_master_participates"], True)
-        self.assertEqual(json.loads(self.prepared_inputs()["image_list_enabled_worker_ids"]), ["worker"])
-        self.assertEqual(result[2], 1)  # Public worker_count still counts remote participants only.
-        self.preflight.assert_awaited_once_with(
-            [{"id": "worker", "name": "Worker", "type": "local", "host": None, "port": 8188}],
-            require_idle=False)
-        self.q.dispatch_worker_prompt.assert_not_awaited()
-        self.assertEqual(self.ps.distributed_pending_jobs, {})
-
-    async def test_orchestrator_setting_excludes_master_processing(self):
-        self.config["settings"]["master_delegate_only"] = True
-        await self.queue()
-        self.assertIs(self.prepared_inputs()["image_list_master_participates"], False)
-        self.assertTrue(self.preflight.call_args.kwargs["require_idle"])
-        # Upstream source stays on master even in orchestrator-only mode.
-        self.assertIn("1", self.q.queue_prompt_payload.call_args.args[0])
-
-    async def test_explicit_override_wins_over_config_in_both_directions(self):
-        for configured, requested in [(True, False), (False, True)]:
-            with self.subTest(configured=configured, requested=requested):
-                self.config["settings"]["master_delegate_only"] = configured
-                await self.queue(delegate_master=requested)
-                self.assertIs(self.prepared_inputs()["image_list_master_participates"], not requested)
-
-    async def test_master_only_with_explicit_empty_worker_selection(self):
-        result = await self.queue(enabled_worker_ids=[])
-        self.assertIs(self.prepared_inputs()["image_list_master_participates"], True)
-        self.assertEqual(json.loads(self.prepared_inputs()["image_list_enabled_worker_ids"]), [])
-        self.assertEqual(result[2], 0)
-        self.preflight.assert_not_awaited()
-
-    async def test_master_only_when_no_workers_are_configured(self):
-        self.config["workers"] = []
-        await self.queue()
-        self.assertIs(self.prepared_inputs()["image_list_master_participates"], True)
-
-    async def test_unavailable_workers_do_not_prevent_enabled_master(self):
-        self.q.select_active_workers.return_value = ([], False)
-        self.q.select_active_workers.side_effect = None
-        self.preflight.return_value = {}
-        await self.queue()
-        self.assertEqual(json.loads(self.prepared_inputs()["image_list_enabled_worker_ids"]), [])
-        self.assertIs(self.prepared_inputs()["image_list_master_participates"], True)
-
-    async def test_busy_workers_do_not_prevent_enabled_master(self):
-        self.preflight.return_value = {}
-        await self.queue()
-        self.assertEqual(json.loads(self.prepared_inputs()["image_list_enabled_worker_ids"]), [])
-
-    async def test_orchestrator_without_workers_fails_before_queue(self):
-        self.config["settings"]["master_delegate_only"] = True
-        with self.assertRaisesRegex(RuntimeError, "orchestrator|Orchestrator"):
-            await self.queue(enabled_worker_ids=[])
-        self.q.queue_prompt_payload.assert_not_awaited()
-
-    async def test_orchestrator_never_inherits_legacy_offline_master_fallback(self):
-        self.config["settings"]["master_delegate_only"] = True
-        self.q.select_active_workers.side_effect = None
-        self.q.select_active_workers.return_value = ([], False)
-        self.preflight.return_value = {}
-        with self.assertRaisesRegex(RuntimeError, "orchestrator|Orchestrator"):
-            await self.queue()
-        self.q.queue_prompt_payload.assert_not_awaited()
-
-    async def test_serverless_only_uses_master_without_provider_dispatch(self):
-        self.config["workers"] = [{"id": "pool", "type": "serverless", "enabled": True}]
-        result = await self.queue()
-        self.assertIs(self.prepared_inputs()["image_list_master_participates"], True)
-        self.assertEqual(result[2], 0)
-        self.q.dispatch_worker_prompt.assert_not_awaited()
-
-    async def test_incompatible_worker_is_not_silently_hidden_by_master(self):
-        self.preflight.side_effect = RuntimeError("upgrade worker")
-        with self.assertRaisesRegex(RuntimeError, "upgrade worker"):
-            await self.queue()
-        self.q.queue_prompt_payload.assert_not_awaited()

tests/test_image_list_runtime_boundaries.py

@@ -1,340 +0,0 @@
-"""Exercise production per-item orchestration through controlled provider seams."""
-import asyncio
-import importlib.util
-import json
-from pathlib import Path
-import sys
-import types
-import unittest
-from unittest.mock import AsyncMock, Mock, patch
-
-import torch
-
-from test_image_list import graph_one_region
-
-ROOT = Path(__file__).resolve().parents[1]
-
-
-class Response:
-    def __init__(self, data=None, status=200):
-        self.data, self.status = data, status
-
-    async def __aenter__(self):
-        return self
-
-    async def __aexit__(self, *_args):
-        return None
-
-    async def json(self):
-        return self.data
-
-    def raise_for_status(self):
-        if self.status >= 400:
-            raise RuntimeError(f'HTTP {self.status}')
-
-
-class RuntimeBoundaryTests(unittest.IsolatedAsyncioTestCase):
-    async def asyncSetUp(self):
-        name = '_image_list_boundaries'
-        self.modules = {}
-        for suffix in ('', '.image_list', '.api', '.api.orchestration', '.utils'):
-            module = types.ModuleType(name + suffix)
-            module.__path__ = []
-            self.modules[name + suffix] = module
-        self.patch_modules = patch.dict(sys.modules, self.modules)
-        self.patch_modules.start()
-        self.addCleanup(self.patch_modules.stop)
-        spec = importlib.util.spec_from_file_location(name + '.image_list.runtime', ROOT / 'image_list/runtime.py')
-        self.rt = importlib.util.module_from_spec(spec)
-        sys.modules[spec.name] = self.rt
-        spec.loader.exec_module(self.rt)
-        self.worker = {'id': 'worker', 'host': 'worker', 'type': 'remote', 'port': 8188}
-        self.config = {'workers': [self.worker], 'settings': {}}
-        self.prompts, self.uploads, self.deletions, self.cancel_requests = [], [], [], []
-        self.reject = False
-        self.interrupted = False
-        self.completed = True
-        self.queue = {'queue_running': [], 'queue_pending': []}
-        self.callback = Mock(side_effect=lambda worker, **_kwargs:
-                             'http://127.0.0.1:8188' if worker['type'] == 'local' else 'http://public-master')
-        self.sync = AsyncMock(side_effect=self.sync_media)
-        self.dispatch = AsyncMock(side_effect=self.dispatch_prompt)
-        self.logs = []
-        dependencies = {
-            '.api.orchestration.dispatch': types.SimpleNamespace(
-                preflight_image_list_workers=AsyncMock(return_value={'worker': True}),
-                dispatch_worker_prompt=self.dispatch),
-            '.api.orchestration.prompt_transform': types.SimpleNamespace(
-                _resolve_direct_workers_for_image_list=lambda config, ids: config['workers']),
-            '.api.orchestration.media_sync': types.SimpleNamespace(
-                sync_worker_media=self.sync, fetch_worker_path_separator=AsyncMock(return_value='/'),
-                convert_paths_for_platform=lambda prompt, separator: prompt),
-            '.utils.config': types.SimpleNamespace(load_config=lambda: self.config,
-                get_worker_timeout_seconds=lambda: 60,
-                get_image_list_job_timeout_seconds=lambda config=None: 0),
-            '.utils.network': types.SimpleNamespace(build_master_url=lambda **kwargs: 'http://public-master',
-                build_master_callback_url=self.callback,
-                build_worker_url=lambda worker, path='': f"http://{worker['id']}{path}",
-                get_client_session=AsyncMock(return_value=self)),
-            '.utils.logging': types.SimpleNamespace(log=self.logs.append),
-        }
-        for suffix, module in dependencies.items():
-            sys.modules[name + suffix] = module
-        self.config_module = dependencies['.utils.config']
-        # Other legacy test modules replace aiohttp during collection. This
-        # controlled provider fixture needs only its explicit timeout/error seam.
-        sys.modules['aiohttp'] = types.SimpleNamespace(
-            ClientTimeout=lambda **kwargs: kwargs, ClientError=ConnectionError)
-        comfy = types.ModuleType('comfy')
-        comfy.model_management = types.SimpleNamespace(processing_interrupted=lambda: self.interrupted)
-        sys.modules['comfy'] = comfy
-        sys.modules['comfy.model_management'] = comfy.model_management
-
-    async def sync_media(self, worker, prompt):
-        prompt['7']['inputs']['image'] = 'shared/ref.png'
-
-    async def dispatch_prompt(self, worker, prompt, meta, **kwargs):
-        if self.reject:
-            raise RuntimeError('prompt rejected')
-        self.prompts.append((json.loads(json.dumps(prompt)), meta))
-        inputs = prompt['5']['inputs']
-        token = inputs['worker_result_token']
-        item = self.rt.consume_worker_item(token)
-        self.rt.set_master_result(self.context, inputs['worker_item_index'], token, {'image': item['image']})
-        return f'prompt-{len(self.prompts)}'
-
-    def post(self, url, *, json=None, data=None, **kwargs):
-        if data is not None:
-            json = __import__('json').loads(data)
-        if url.endswith('/distributed/image-list/item'):
-            self.uploads.append(json['token'])
-            self.rt.store_worker_item(json['token'], json['item'])
-        else:
-            self.cancel_requests.append((url, json))
-        return Response({})
-
-    def delete(self, url, **kwargs):
-        token = url.rsplit('/', 1)[-1]
-        self.deletions.append(token)
-        self.rt.delete_worker_item(token)
-        return Response({})
-
-    def get(self, url, **kwargs):
-        if url.endswith('/queue'):
-            return Response(self.queue)
-        pid = url.rsplit('/', 1)[-1]
-        return Response({pid: {'status': {'completed': True, 'status_str': 'success'}}}
-                        if self.completed else {})
-
-    def create(self, count=2):
-        prompt = graph_one_region()
-        prompt['7'] = {'class_type': 'LoadImage', 'inputs': {'image': 'shared/ref.png [output]'}}
-        prompt['4']['inputs']['reference'] = ['7', 0]
-        self.metadata = {'nodes': [{'id': 42}], 'extra': {'label': 'generic fixture'}}
-        self.context = self.rt.create_master_job('3', '5', prompt,
-            [{'image': torch.zeros((1, 2, 2, 3)), 'index': i} for i in range(count)],
-            ['worker'], 'http://public-master', self.metadata)
-        return self.context
-
-    async def test_each_assignment_reuses_media_paths_with_fresh_tokens_and_metadata(self):
-        context = self.create()
-        await self.rt.run_image_list_job(context)
-        self.assertEqual(self.sync.await_count, 1)
-        self.assertEqual([p['7']['inputs']['image'] for p, _ in self.prompts],
-                         ['shared/ref.png', 'shared/ref.png'])
-        tokens = [p['5']['inputs']['worker_result_token'] for p, _ in self.prompts]
-        self.assertEqual(len(set(tokens)), 2)
-        self.assertEqual([p['5']['inputs']['worker_item_index'] for p, _ in self.prompts], [0, 1])
-        self.assertEqual([meta for _, meta in self.prompts], [self.metadata, self.metadata])
-
-    async def test_annotated_media_roots_survive_windows_prep_and_repeated_items(self):
-        import tempfile
-
-        context = self.create(2)
-        state = self.rt.get_master_job(context)
-        state.original_prompt['8'] = {'class_type': 'LoadImage', 'inputs': {'image': 'shared\\ref.png [temp]'}}
-        state.original_prompt['4']['inputs']['second_reference'] = ['8', 0]
-        name = '_image_list_boundaries.api.orchestration.media_sync'
-        spec = importlib.util.spec_from_file_location(name, ROOT / 'api/orchestration/media_sync.py')
-        assert spec is not None and spec.loader is not None
-        media = importlib.util.module_from_spec(spec)
-        sys.modules['_image_list_boundaries.utils.logging'].debug_log = self.logs.append
-        sys.modules['_image_list_boundaries.utils.trace_logger'] = types.SimpleNamespace(
-            trace_debug=lambda *args: None, trace_info=lambda *args: None)
-        spec.loader.exec_module(media)
-        media.fetch_worker_path_separator = AsyncMock(return_value='\\')
-        converted = Mock(wraps=media.convert_paths_for_platform)
-        media.convert_paths_for_platform = converted
-        uploads = []
-
-        async def upload(worker, filename, content, digest, mime, **kwargs):
-            uploads.append((filename, content))
-            return True, f'worker-input/{len(uploads)}.png'
-
-        media._upload_media_to_worker = AsyncMock(side_effect=upload)
-        media.sync_worker_media = AsyncMock(wraps=media.sync_worker_media)
-        with tempfile.TemporaryDirectory(prefix='image-list-media-') as directory:
-            root = Path(directory)
-            for kind in ('input', 'output', 'temp'):
-                path = root / kind / 'shared/ref.png'
-                path.parent.mkdir(parents=True)
-                path.write_bytes(kind.encode())
-
-            def resolve_annotated(filename):
-                for kind in ('output', 'temp', 'input'):
-                    suffix = f' [{kind}]'
-                    if filename.endswith(suffix):
-                        return str(root / kind / filename[:-len(suffix)])
-                return str(root / 'input' / filename)
-
-            with patch.dict(sys.modules, {name: media, 'folder_paths': types.SimpleNamespace(
-                    get_annotated_filepath=resolve_annotated)}):
-                await self.rt.run_image_list_job(context)
-        self.assertEqual(uploads, [('shared/ref.png [output]', b'output'),
-                                   ('shared/ref.png [temp]', b'temp')])
-        media.sync_worker_media.assert_awaited_once()
-        media.fetch_worker_path_separator.assert_awaited_once()
-        converted.assert_called_once()
-        self.assertEqual(converted.call_args.args[1], '\\')
-        for prompt, _ in self.prompts:
-            self.assertEqual(prompt['7']['inputs']['image'], 'worker-input/1.png')
-            self.assertEqual(prompt['8']['inputs']['image'], 'worker-input/2.png')
-        self.assertEqual(state.original_prompt['7']['inputs']['image'], 'shared/ref.png [output]')
-
-    async def test_local_worker_callback_does_not_use_public_master_address(self):
-        self.worker['type'] = 'local'
-        await self.rt.run_image_list_job(self.create(1))
-        self.assertEqual(self.prompts[0][0]['5']['inputs']['master_callback_url'], 'http://127.0.0.1:8188')
-        self.callback.assert_called()
-
-    async def test_remote_worker_uses_topology_callback_helper(self):
-        await self.rt.run_image_list_job(self.create(1))
-        self.assertEqual(self.prompts[0][0]['5']['inputs']['master_callback_url'], 'http://public-master')
-        self.callback.assert_called()
-
-    def test_metadata_is_forwarded_into_master_start_prompt(self):
-        from test_prompt_transform import pt
-        prompt = graph_one_region()
-        metadata = {'nodes': [{'id': 1}]}
-        index = pt.PromptIndex(prompt)
-        transformed = pt.apply_participant_overrides(
-            prompt, 'master', ['worker'], {'3': 'run'}, 'http://public-master', False,
-            index, workflow_meta=metadata)
-        self.assertEqual(transformed['3']['inputs']['image_list_workflow_meta'], metadata)
-
-    async def test_aggregate_input_limit_fails_before_upload(self):
-        self.rt.MAX_REQUEST_BYTES = 64
-        with self.assertRaisesRegex(ValueError, 'complete.*request'):
-            await self.rt.run_image_list_job(self.create(1))
-        self.assertFalse(self.uploads)
-        self.assertFalse(self.rt._WORKER_ITEMS)
-
-    def test_tensor_and_combined_request_limits_preserve_dtype(self):
-        item = {'image': torch.zeros((1, 4, 4, 3)), 'index': 0}
-        encoded = self.rt.encode_item(item)
-        payload = {'context_id': 'example', 'token': 't' * 48, 'item': encoded}
-        single = self.rt.encode_request_payload(payload)
-        self.assertEqual(json.loads(single)['item']['image']['dtype'], 'float32')
-        self.rt.MAX_REQUEST_BYTES = len(single)
-        self.assertEqual(self.rt.encode_request_payload(payload), single)
-        item['values'] = {'value_1': item['image']}
-        item['mask'] = torch.zeros((1, 4, 4))
-        payload['item'] = self.rt.encode_item(item)
-        with self.assertRaisesRegex(ValueError, 'complete.*request'):
-            self.rt.encode_request_payload(payload)
-        # Size is checked before allocating/copying a 4K tensor onto the CPU.
-        with self.assertRaisesRegex(ValueError, 'transport limit'):
-            self.rt.encode_value(torch.empty((1, 2160, 3840, 3), device='meta'), 'IMAGE')
-
-    async def test_rejected_submission_reclaims_uploaded_token_without_prompt_id(self):
-        self.reject = True
-        with self.assertRaisesRegex(RuntimeError, 'prompt rejected'):
-            await self.rt.run_image_list_job(self.create(1))
-        self.assertEqual(self.deletions, self.uploads)
-        self.assertFalse(self.rt._WORKER_ITEMS)
-        self.assertFalse(self.rt._MASTER_JOBS)
-        self.assertFalse(self.cancel_requests)
-
-    async def test_user_interrupt_during_pending_submission_is_bounded(self):
-        entered = asyncio.Event()
-        async def never_ack(*args, **kwargs):
-            entered.set()
-            await asyncio.Event().wait()
-        self.dispatch.side_effect = never_ack
-        self.rt.SUBMISSION_CANCEL_GRACE_SECONDS = 0.02
-        task = asyncio.create_task(self.rt.run_image_list_job(self.create(2)))
-        await asyncio.wait_for(entered.wait(), 1)
-        self.interrupted = True
-        with self.assertRaises(self.rt.ImageListCancelledError):
-            await asyncio.wait_for(task, 0.8)
-        self.assertEqual(len(self.uploads), 1)
-        self.assertEqual(self.deletions, self.uploads)
-        self.assertFalse(self.rt._MASTER_JOBS)
-        self.assertTrue(any('ambiguous' in line for line in self.logs), self.logs)
-
-    async def test_owned_running_prompt_survives_cold_load_and_old_hour_cap(self):
-        self.completed = False
-        context = self.create(2)
-        clock = [0.0]
-        calls = [0]
-        original_get = self.get
-        def get(url, **kwargs):
-            if '/history/' in url:
-                calls[0] += 1
-                clock[0] += 3700
-                self.completed = calls[0] >= 3
-            if url.endswith('/queue'):
-                return Response({'queue_running': [[0, 'prompt-1']] if calls[0] == 2 else [],
-                                 'queue_pending': [[0, 'prompt-1']] if calls[0] == 1 else []})
-            return original_get(url, **kwargs)
-        self.rt.time = types.SimpleNamespace(monotonic=lambda: clock[0])
-        with patch.object(self, 'get', side_effect=get):
-            result = await self.rt.run_image_list_job(context)
-        self.assertEqual(len(result[0]), 2)
-        self.assertGreaterEqual(clock[0], 7200)
-        self.assertFalse(self.cancel_requests)
-
-    async def test_unrelated_busy_prompt_does_not_refresh_owned_liveness(self):
-        self.completed = False
-        clock = [0.0]
-        original_get = self.get
-        def get(url, **kwargs):
-            if '/history/' in url:
-                clock[0] += 70
-            if url.endswith('/queue'):
-                return Response({'queue_running': [[0, 'not-our-prompt']], 'queue_pending': []})
-            return original_get(url, **kwargs)
-        self.rt.time = types.SimpleNamespace(monotonic=lambda: clock[0])
-        with patch.object(self, 'get', side_effect=get):
-            with self.assertRaisesRegex(TimeoutError, 'liveness'):
-                await self.rt.run_image_list_job(self.create(1))
-        self.assertEqual(self.cancel_requests, [
-            ('http://worker/queue', {'delete': ['prompt-1']}),
-            ('http://worker/interrupt', {'prompt_id': 'prompt-1'})])
-
-    async def test_failed_queue_cleanup_still_interrupts_and_reclaims_staging(self):
-        self.completed = False
-        self.config_module.get_image_list_job_timeout_seconds = lambda config=None: 0.03
-        original_post = self.post
-        def post(url, **kwargs):
-            result = original_post(url, **kwargs)
-            return Response(status=500) if url.endswith('/queue') else result
-        with patch.object(self, 'post', side_effect=post):
-            with self.assertRaisesRegex(TimeoutError, 'total deadline'):
-                await self.rt.run_image_list_job(self.create(2))
-        self.assertEqual(self.cancel_requests, [
-            ('http://worker/queue', {'delete': ['prompt-1']}),
-            ('http://worker/interrupt', {'prompt_id': 'prompt-1'})])
-        self.assertEqual(self.deletions, self.uploads)
-        self.assertTrue(any('cleanup unconfirmed' in line for line in self.logs), self.logs)
-        self.assertEqual(self.dispatch.await_count, 1)
-
-    async def test_configured_total_deadline_cleans_owned_work(self):
-        self.completed = False
-        self.config_module.get_image_list_job_timeout_seconds = lambda config=None: 0.03
-        self.queue = {'queue_running': [[0, 'prompt-1']], 'queue_pending': []}
-        with self.assertRaisesRegex(TimeoutError, 'total deadline'):
-            await asyncio.wait_for(self.rt.run_image_list_job(self.create(2)), 1)
-        self.assertEqual(len(self.uploads), 1)
-        self.assertEqual(self.deletions, self.uploads)
-        self.assertFalse(self.rt._MASTER_JOBS)

tests/test_image_list_timeout_config.py

@@ -1,106 +0,0 @@
-import importlib.util
-import sys
-import types
-from pathlib import Path
-
-import pytest
-
-
-def _load_config_module():
-    module_path = Path(__file__).resolve().parents[1] / "utils" / "config.py"
-    package_name = "image_list_timeout_config_testpkg"
-
-    for mod_name in list(sys.modules):
-        if mod_name == package_name or mod_name.startswith(f"{package_name}."):
-            del sys.modules[mod_name]
-
-    root_pkg = types.ModuleType(package_name)
-    root_pkg.__path__ = [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.HEARTBEAT_TIMEOUT = 30
-    sys.modules[f"{package_name}.constants"] = constants_module
-
-    spec = importlib.util.spec_from_file_location(f"{package_name}.config", 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
-
-
-config_module = _load_config_module()
-
-
-SETTING_KEY = "image_list_job_timeout_seconds"
-
-
-def _config_with_timeout(value):
-    return {"settings": {SETTING_KEY: value}}
-
-
-def test_default_config_sets_image_list_job_timeout_to_zero():
-    assert config_module.get_default_config()["settings"][SETTING_KEY] == 0.0
-
-
-def test_missing_value_defaults_to_disabled_timeout():
-    assert config_module.get_image_list_job_timeout_seconds({}) == 0.0
-    assert config_module.get_image_list_job_timeout_seconds({"settings": {}}) == 0.0
-
-
-def test_none_config_loads_current_config(monkeypatch):
-    cfg = _config_with_timeout("12.5")
-    monkeypatch.setattr(config_module, "load_config", lambda: cfg)
-
-    assert config_module.get_image_list_job_timeout_seconds() == 12.5
-
-
-@pytest.mark.parametrize(
-    ("value", "expected"),
-    [
-        (0, 0.0),
-        (0.0, 0.0),
-        ("0", 0.0),
-        (" 0.0 ", 0.0),
-        (1, 1.0),
-        (2.5, 2.5),
-        ("3", 3.0),
-        (" 4.25 ", 4.25),
-    ],
-)
-def test_accepts_zero_and_positive_numeric_values(value, expected):
-    result = config_module.get_image_list_job_timeout_seconds(_config_with_timeout(value))
-
-    assert result == expected
-    assert isinstance(result, float)
-
-
-@pytest.mark.parametrize(
-    "value",
-    [
-        True,
-        False,
-        -1,
-        -0.01,
-        float("nan"),
-        float("inf"),
-        float("-inf"),
-        "not-a-number",
-        "",
-        "   ",
-        "-1",
-        "nan",
-        "inf",
-        None,
-        [],
-        {},
-    ],
-)
-def test_rejects_invalid_values_with_setting_name(value):
-    with pytest.raises(ValueError, match=SETTING_KEY):
-        config_module.get_image_list_job_timeout_seconds(_config_with_timeout(value))

tests/test_image_list_websocket_deadline.py

@@ -1,212 +0,0 @@
-import asyncio
-import importlib.util
-import sys
-import types
-import unittest
-import uuid
-from pathlib import Path
-from typing import Any, cast
-from unittest.mock import patch
-
-import aiohttp
-from aiohttp import web
-
-
-def _load_dispatch_module():
-    module_path = Path(__file__).resolve().parents[1] / "api" / "orchestration" / "dispatch.py"
-    package_name = f"image_list_ws_deadline_testpkg_{uuid.uuid4().hex}"
-    package_modules = {}
-
-    for suffix in ("", ".api", ".api.orchestration", ".utils"):
-        package = types.ModuleType(package_name + suffix)
-        package.__path__ = []
-        package_modules[package_name + suffix] = package
-
-    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)
-    package_modules[logging_module.__name__] = logging_module
-
-    network_module = types.ModuleType(f"{package_name}.utils.network")
-    setattr(network_module, "build_worker_url", lambda *_args, **_kwargs: "http://example.invalid")
-
-    async def _unused_probe_worker(*_args, **_kwargs):
-        return None
-
-    async def _unused_get_client_session(*_args, **_kwargs):
-        raise RuntimeError("get_client_session must be patched by this test")
-
-    setattr(network_module, "probe_worker", _unused_probe_worker)
-    setattr(network_module, "get_client_session", _unused_get_client_session)
-    package_modules[network_module.__name__] = network_module
-
-    module_name = f"{package_name}.api.orchestration.dispatch"
-    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)
-    package_modules[module_name] = module
-
-    previous_modules = {name: sys.modules.get(name) for name in package_modules}
-    try:
-        sys.modules.update(package_modules)
-        spec.loader.exec_module(module)
-    finally:
-        for name, previous in previous_modules.items():
-            if previous is None:
-                sys.modules.pop(name, None)
-            else:
-                sys.modules[name] = previous
-    return module
-
-
-dispatch = _load_dispatch_module()
-
-
-class WebsocketDispatchDeadlineTests(unittest.IsolatedAsyncioTestCase):
-    async def asyncSetUp(self):
-        self.mode = None
-        self.received_payload = None
-        self.dispatch_received = asyncio.Event()
-        self.peer_finished = asyncio.Event()
-        self.irrelevant_messages_sent = 0
-
-        app = web.Application()
-        app.router.add_get("/distributed/worker_ws", self._handle_websocket)
-        self.runner = web.AppRunner(app)
-        await self.runner.setup()
-        self.site = web.TCPSite(self.runner, "127.0.0.1", 0)
-        await self.site.start()
-        server = cast(Any, self.site._server)
-        port = server.sockets[0].getsockname()[1]
-        self.worker_url = f"http://127.0.0.1:{port}"
-        self.session = aiohttp.ClientSession()
-
-    async def asyncTearDown(self):
-        await self.session.close()
-        await self.runner.cleanup()
-
-    async def _get_client_session(self):
-        return self.session
-
-    async def _handle_websocket(self, request):
-        ws = web.WebSocketResponse()
-        await ws.prepare(request)
-        try:
-            payload = await ws.receive_json()
-            self.received_payload = payload
-            self.dispatch_received.set()
-
-            if self.mode in ("success", "slow-close"):
-                await ws.send_json({
-                    "type": "dispatch_ack",
-                    "request_id": "unrelated-request",
-                    "ok": True,
-                    "prompt_id": "wrong-prompt",
-                })
-                await ws.send_json({"type": "heartbeat"})
-                await ws.send_json({
-                    "type": "dispatch_ack",
-                    "request_id": payload["request_id"],
-                    "ok": True,
-                    "prompt_id": "prompt-123",
-                })
-                if self.mode == "slow-close":
-                    await asyncio.sleep(0.15)
-            elif self.mode == "rejection":
-                await ws.send_json({
-                    "type": "dispatch_ack",
-                    "request_id": payload["request_id"],
-                    "ok": False,
-                    "error": "Invalid prompt",
-                    "validation_error": {"node": "bad"},
-                    "node_errors": {"bad": "missing input"},
-                })
-            elif self.mode in ("no-ack", "wait-for-cancel"):
-                async def send_irrelevant_messages():
-                    while True:
-                        try:
-                            await ws.send_json({"type": "heartbeat"})
-                            self.irrelevant_messages_sent += 1
-                            await ws.send_json({
-                                "type": "dispatch_ack",
-                                "request_id": "unrelated-request",
-                                "ok": True,
-                                "prompt_id": "not-this-request",
-                            })
-                            self.irrelevant_messages_sent += 1
-                        except (ConnectionError, RuntimeError):
-                            return
-                        await asyncio.sleep(0.01)
-
-                sender = asyncio.create_task(send_irrelevant_messages())
-                try:
-                    async for _message in ws:
-                        pass
-                finally:
-                    sender.cancel()
-                    try:
-                        await sender
-                    except asyncio.CancelledError:
-                        pass
-        finally:
-            self.peer_finished.set()
-        return ws
-
-    async def _dispatch(self, timeout):
-        with patch.object(dispatch, "get_client_session", self._get_client_session):
-            return await dispatch._dispatch_via_websocket(
-                self.worker_url,
-                {"prompt": {"node": {"class_type": "Example"}}, "workflow": {"name": "test"}},
-                "client-456",
-                timeout=timeout,
-            )
-
-    async def test_returns_matching_acknowledgement_prompt_id(self):
-        self.mode = "success"
-
-        prompt_id = await self._dispatch(timeout=0.5)
-
-        self.assertEqual(prompt_id, "prompt-123")
-        received_payload = self.received_payload
-        assert received_payload is not None
-        self.assertEqual(received_payload["type"], "dispatch_prompt")
-        self.assertTrue(received_payload["request_id"])
-        self.assertEqual(received_payload["client_id"], "client-456")
-        self.assertEqual(received_payload["prompt"], {"node": {"class_type": "Example"}})
-        self.assertEqual(received_payload["workflow"], {"name": "test"})
-
-    async def test_acknowledged_prompt_id_survives_slow_close(self):
-        self.mode = "slow-close"
-        self.assertEqual(await self._dispatch(timeout=0.03), "prompt-123")
-
-    async def test_matching_rejection_preserves_worker_error_details(self):
-        self.mode = "rejection"
-
-        with self.assertRaisesRegex(RuntimeError, "Invalid prompt.*validation_error=.*node_errors="):
-            await self._dispatch(timeout=0.5)
-
-    async def test_repeated_irrelevant_messages_do_not_reset_ack_deadline(self):
-        self.mode = "no-ack"
-        started = asyncio.get_running_loop().time()
-
-        with self.assertRaisesRegex(RuntimeError, "submission status is ambiguous"):
-            await asyncio.wait_for(self._dispatch(timeout=0.12), timeout=0.5)
-
-        elapsed = asyncio.get_running_loop().time() - started
-        self.assertLess(elapsed, 0.5)
-        self.assertGreaterEqual(self.irrelevant_messages_sent, 4)
-        await asyncio.wait_for(self.peer_finished.wait(), timeout=0.5)
-
-    async def test_outer_cancellation_propagates_and_closes_peer(self):
-        self.mode = "wait-for-cancel"
-        task = asyncio.create_task(self._dispatch(timeout=2.0))
-        await asyncio.wait_for(self.dispatch_received.wait(), timeout=0.5)
-
-        task.cancel()
-        with self.assertRaises(asyncio.CancelledError):
-            await asyncio.wait_for(task, timeout=0.5)
-        await asyncio.wait_for(self.peer_finished.wait(), timeout=0.5)
-
-
-if __name__ == "__main__":
-    unittest.main()

tests/test_payload_parsers.py

@@ -86,25 +86,27 @@ class ParseTilesFromFormTests(unittest.TestCase):
 
     # --- happy paths ---
 
-    def test_single_tile_returns_image_and_metadata(self):
+    def test_single_tile_returns_one_entry(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)
@@ -136,6 +138,16 @@ 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_prompt_transform.py

@@ -260,40 +260,6 @@ class PrunePromptForWorkerTests(unittest.TestCase):
         preview_nodes = [n for n in result.values() if n.get("class_type") == "PreviewImage"]
         self.assertEqual(len(preview_nodes), 0)
 
-    def test_worker_sink_replaces_preview_image(self):
-        result = pt.prune_prompt_for_worker(_linear_prompt(), use_worker_sink=True)
-        sinks = [n for n in result.values() if n.get("class_type") == pt.WORKER_SINK_CLASS]
-        self.assertEqual(len(sinks), 1)
-        self.assertEqual(sinks[0]["inputs"], {"value": ["4", 0]})
-        self.assertFalse(any(n.get("class_type") in ("PreviewImage", "PreviewAudio") for n in result.values()))
-
-    def test_worker_sink_uses_audio_output_for_audio_only_collector(self):
-        result = pt.prune_prompt_for_worker(_audio_only_prompt(), use_worker_sink=True)
-        sinks = [n for n in result.values() if n.get("class_type") == pt.WORKER_SINK_CLASS]
-        self.assertEqual([s["inputs"] for s in sinks], [{"value": ["2", 1]}])
-        self.assertFalse(any(n.get("class_type") == "PreviewAudio" for n in result.values()))
-
-    def test_worker_sink_for_distributed_upscale(self):
-        prompt = {
-            "1": {"class_type": "LoadImage", "inputs": {}},
-            "2": {"class_type": "UltimateSDUpscaleDistributed", "inputs": {"upscaled_image": ["1", 0]}},
-            "3": {"class_type": "SaveImage", "inputs": {"images": ["2", 0]}},
-        }
-        result = pt.prune_prompt_for_worker(prompt, use_worker_sink=True)
-        sinks = [n for n in result.values() if n.get("class_type") == pt.WORKER_SINK_CLASS]
-        self.assertEqual([s["inputs"] for s in sinks], [{"value": ["2", 0]}])
-        self.assertNotIn("3", result)
-
-    def test_worker_sink_ids_do_not_collide_with_original_prompt(self):
-        prompt = _linear_prompt()
-        result = pt.prune_prompt_for_worker(prompt, use_worker_sink=True)
-        sink_ids = [k for k, n in result.items() if n.get("class_type") == pt.WORKER_SINK_CLASS]
-        self.assertTrue(all(k not in prompt for k in sink_ids))
-
-    def test_no_worker_sink_when_no_downstream(self):
-        result = pt.prune_prompt_for_worker(_collector_only_prompt(), use_worker_sink=True)
-        self.assertFalse(any(n.get("class_type") == pt.WORKER_SINK_CLASS for n in result.values()))
-
     def test_unrelated_nodes_are_pruned(self):
         prompt = {
             "1": {"class_type": "DistributedCollector", "inputs": {}},
@@ -336,8 +302,8 @@ class PrepareDelegateMasterPromptTests(unittest.TestCase):
         self.assertNotIn("1", result)
         self.assertNotIn("2", result)
 
-    def test_replaces_dangling_upstream_ref_with_one_empty_image_placeholder(self):
-        """Collector must point to exactly one valid placeholder after pruning."""
+    def test_removes_dangling_upstream_refs(self):
+        """Collector must not retain dangling refs to pruned upstream nodes."""
         prompt = _delegate_prompt()
         result = pt.prepare_delegate_master_prompt(prompt, ["3"])
         collector_inputs = result["3"].get("inputs", {})
@@ -348,6 +314,10 @@ class PrepareDelegateMasterPromptTests(unittest.TestCase):
         self.assertNotEqual(source_id, "2")
         self.assertIn(source_id, result)
         self.assertEqual(result[source_id].get("class_type"), "DistributedEmptyImage")
+
+    def test_injects_empty_image_placeholder(self):
+        prompt = _delegate_prompt()
+        result = pt.prepare_delegate_master_prompt(prompt, ["3"])
         empty_nodes = [(nid, n) for nid, n in result.items() if n.get("class_type") == "DistributedEmptyImage"]
         self.assertEqual(len(empty_nodes), 1)
         placeholder_id = empty_nodes[0][0]
@@ -414,6 +384,13 @@ class PrepareDelegateMasterPromptTests(unittest.TestCase):
         self.assertEqual(result["15"], prompt["15"])
         self.assertEqual(result["9"]["inputs"]["filename_prefix"], ["15", 0])
 
+    def test_does_not_preserve_non_primitive_upstream_for_collector(self):
+        prompt = _delegate_prompt()
+        result = pt.prepare_delegate_master_prompt(prompt, ["3"])
+
+        self.assertNotIn("2", result)
+        self.assertNotEqual(result["3"]["inputs"]["images"], ["2", 0])
+
     def test_preserves_load_image_for_switch_alternate_required_input(self):
         """Delegate-only master keeps LoadImage inputs needed by switches."""
         prompt = {

tests/test_runpod_client.py

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

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

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

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

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

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

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

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

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

@@ -1,98 +0,0 @@
-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_transfer_codec.py

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

@@ -1,126 +0,0 @@
-"""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(f"{PKG}.nodes.image_list", DistributedImageListStart=object, DistributedImageListEnd=object)
-    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

@@ -1,37 +0,0 @@
-"""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
-
-
-def test_image_list_network_runtime():
-    comfy_root = os.environ.get('COMFYUI_SOURCE_ROOT')
-    if not comfy_root:
-        pytest.skip('Set COMFYUI_SOURCE_ROOT for model-free multi-process HTTP acceptance')
-    helper = Path(__file__).parent / 'helpers' / 'image_list_network_check.py'
-    result = subprocess.run(
-        [sys.executable, str(helper), str(Path(comfy_root).resolve())],
-        capture_output=True, text=True, timeout=180,
-        env={**os.environ, 'COMFYUI_IS_WORKER': '1'},
-    )
-    assert result.returncode == 0, result.stdout + '\n' + result.stderr
-    assert 'IMAGE_LIST_NETWORK_OK' in result.stdout
-    assert 'IMAGE_LIST_CANCEL_OK' in result.stdout
-    assert 'IMAGE_LIST_MASTER_OK' in result.stdout

tests/test_worker_ports.py

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

@@ -1,16 +1,11 @@
 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
@@ -26,7 +21,7 @@ def _load_process_module(module_filename: str):
     sys.modules[package_name] = root_pkg
 
     workers_pkg = types.ModuleType(f"{package_name}.workers")
-    workers_pkg.__path__ = [str(module_path.parents[1])]
+    workers_pkg.__path__ = []
     sys.modules[f"{package_name}.workers"] = workers_pkg
 
     process_pkg = types.ModuleType(f"{package_name}.workers.process")
@@ -44,23 +39,8 @@ def _load_process_module(module_filename: str):
 
     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,
@@ -90,145 +70,6 @@ class ComfyRootDiscoveryTests(unittest.TestCase):
 
 
 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(
@@ -289,64 +130,5 @@ class LaunchCommandBuilderTests(unittest.TestCase):
         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()

utils/config.py

@@ -2,13 +2,10 @@
 Configuration management for ComfyUI-Distributed.
 """
 import asyncio
-import copy
 import os
 import json
-import math
 from contextlib import asynccontextmanager
 from .logging import log
-from .secret_fields import is_secret_field_name as _is_secret_field_name
 
 # Import defaults for timeout fallbacks
 from .constants import HEARTBEAT_TIMEOUT
@@ -18,52 +15,6 @@ _config_cache = None
 _config_mtime = 0.0
 _config_lock = asyncio.Lock()
 
-_SERVERLESS_OBSOLETE_WORKER_FIELDS = {"host", "port", "listen_port", "cuda_device", "extra_args"}
-
-
-def _contains_serverless_secret(value):
-    if isinstance(value, dict):
-        return any(
-            (isinstance(key, str) and _is_secret_field_name(key))
-            or _contains_serverless_secret(item)
-            for key, item in value.items()
-        )
-    if isinstance(value, (list, tuple)):
-        return any(_contains_serverless_secret(item) for item in value)
-    return False
-
-
-def _normalize_serverless_workers(config):
-    """Remove obsolete local-worker fields without adding local defaults."""
-    if not isinstance(config, dict):
-        return config
-    workers = config.get("workers")
-    if isinstance(workers, list):
-        for worker in workers:
-            if isinstance(worker, dict) and worker.get("type") == "serverless":
-                for field in _SERVERLESS_OBSOLETE_WORKER_FIELDS:
-                    worker.pop(field, None)
-    return config
-
-
-def _validated_config_for_save(config):
-    """Return a safe normalized copy ready for gpu_config.json."""
-    if _contains_serverless_secret(config):
-        raise ValueError("API keys must be stored outside the config file")
-    saved = _normalize_serverless_workers(copy.deepcopy(config))
-    workers = saved.get("workers") if isinstance(saved, dict) else None
-    if isinstance(workers, list) and any(
-        isinstance(worker, dict) and worker.get("type") == "serverless" for worker in workers
-    ):
-        from .serverless_config import validate_serverless_worker
-        for index, worker in enumerate(workers):
-            if isinstance(worker, dict) and worker.get("type") == "serverless":
-                saved_worker = validate_serverless_worker(worker)
-                for field in _SERVERLESS_OBSOLETE_WORKER_FIELDS:
-                    saved_worker.pop(field, None)
-                workers[index] = saved_worker
-    return saved
-
 
 def _config_path():
     return CONFIG_FILE
@@ -82,8 +33,7 @@ def get_default_config():
             "worker_probe_concurrency": 8,
             "worker_prep_concurrency": 4,
             "media_sync_concurrency": 2,
-            "media_sync_timeout_seconds": 120,
-            "image_list_job_timeout_seconds": 0.0
+            "media_sync_timeout_seconds": 120
         },
         "tunnel": {
             "status": "stopped",
@@ -112,7 +62,7 @@ def _merge_with_defaults(data, defaults):
         if key not in merged:
             merged[key] = value
 
-    return _normalize_serverless_workers(merged)
+    return merged
 
 
 def invalidate_config_cache():
@@ -147,12 +97,11 @@ def load_config():
     return _config_cache
 
 def save_config(config):
-    """Saves normalized configuration without allowing API keys in the config file."""
+    """Saves the configuration to file."""
     tmp_path = f"{_config_path()}.tmp"
     try:
-        saved_config = _validated_config_for_save(config)
         with open(tmp_path, 'w', encoding='utf-8') as f:
-            json.dump(saved_config, f, indent=2)
+            json.dump(config, f, indent=2)
             f.flush()
             os.fsync(f.fileno())
         os.replace(tmp_path, _config_path())
@@ -172,21 +121,12 @@ async def config_transaction():
     """Acquire config lock, yield loaded config, and save if changed."""
     async with _config_lock:
         config = load_config()
-        original_config = copy.deepcopy(config)
         original_snapshot = json.dumps(config, sort_keys=True)
-        try:
-            yield config
-        except BaseException:
-            if isinstance(config, dict):
-                config.clear()
-                config.update(original_config)
-            raise
+        yield config
         updated_snapshot = json.dumps(config, sort_keys=True)
-        if updated_snapshot != original_snapshot and not save_config(config):
-            if isinstance(config, dict):
-                config.clear()
-                config.update(original_config)
-            raise RuntimeError("Failed to save config")
+        if updated_snapshot != original_snapshot:
+            if not save_config(config):
+                raise RuntimeError("Failed to save config")
 
 def ensure_config_exists():
     """Creates default config file if it doesn't exist. Used by __init__.py"""
@@ -217,37 +157,6 @@ def get_worker_timeout_seconds(default: int = HEARTBEAT_TIMEOUT) -> int:
         return max(1, int(default))
 
 
-def get_image_list_job_timeout_seconds(config=None) -> float:
-    """Return the validated image-list total-job timeout in seconds.
-
-    A value of zero disables the total-job deadline. Missing settings default
-    to zero; configured values must be finite, non-negative numbers.
-    """
-    setting_key = "settings.image_list_job_timeout_seconds"
-    cfg = load_config() if config is None else config
-    if not isinstance(cfg, dict):
-        raise ValueError(f"{setting_key} must be a finite non-negative number")
-
-    settings = cfg.get("settings", {})
-    if not isinstance(settings, dict):
-        raise ValueError(f"{setting_key} must be a finite non-negative number")
-
-    value = settings.get("image_list_job_timeout_seconds", 0.0)
-    if isinstance(value, bool) or not isinstance(value, (int, float, str)):
-        raise ValueError(f"{setting_key} must be a finite non-negative number")
-    if isinstance(value, str) and not value.strip():
-        raise ValueError(f"{setting_key} must be a finite non-negative number")
-
-    try:
-        timeout = float(value)
-    except (OverflowError, ValueError):
-        raise ValueError(f"{setting_key} must be a finite non-negative number") from None
-
-    if not math.isfinite(timeout) or timeout < 0:
-        raise ValueError(f"{setting_key} must be a finite non-negative number")
-    return timeout
-
-
 def is_master_delegate_only() -> bool:
     """Returns True when master should skip local workload and act as orchestrator only."""
     try:

utils/constants.py

@@ -42,17 +42,6 @@ MEMORY_CLEAR_DELAY = 0.5
 # Batch processing
 MAX_BATCH = int(os.environ.get('COMFYUI_MAX_BATCH', '20'))  # Maximum items per batch to prevent timeouts/OOM (~100MB chunks for 512x512 PNGs)
 
-# Compressed worker -> master result transfer (DistributedCollector)
-# Hard cap on one multipart request body. Must not exceed the master's
-# --max-upload-size (ComfyUI default: 100 MiB), which rejects larger requests.
-TRANSFER_MAX_REQUEST_BYTES = int(os.environ.get('COMFYUI_DISTRIBUTED_TRANSFER_MAX_REQUEST_BYTES', str(96 * 1024 * 1024)))
-# Target media bytes per request when batching several frames into one
-# request; a single frame may exceed this up to the hard cap above.
-TRANSFER_MAX_CHUNK_BYTES = min(
-    int(os.environ.get('COMFYUI_DISTRIBUTED_TRANSFER_CHUNK_BYTES', str(32 * 1024 * 1024))),
-    TRANSFER_MAX_REQUEST_BYTES,
-)
-
 # Heartbeat monitoring
 HEARTBEAT_INTERVAL = float(os.environ.get('COMFYUI_HEARTBEAT_INTERVAL', '10'))  # Heartbeat/check interval in seconds
 HEARTBEAT_TIMEOUT = int(os.environ.get('COMFYUI_HEARTBEAT_TIMEOUT', '60'))  # Worker heartbeat timeout in seconds (default 60s)

utils/network.py

@@ -87,8 +87,6 @@ def _split_host_and_port(host):
 
 def build_worker_url(worker, endpoint=""):
     """Construct the worker base URL with optional endpoint."""
-    if worker.get("type") == "serverless":
-        raise ValueError("Serverless workers have no ComfyUI worker URL")
     host = (worker.get("host") or "").strip()
     port = int(worker.get("port", worker.get("listen_port", 8188)) or 8188)
 

utils/secret_fields.py

@@ -1,9 +0,0 @@
-"""Shared classification for credential-bearing configuration fields."""
-
-
-def is_secret_field_name(key: object) -> bool:
-    """Return whether a field name identifies an API key or credential."""
-    if not isinstance(key, str):
-        return False
-    normalized = "".join(character for character in key.casefold() if character.isalnum())
-    return normalized in {"apikey", "clearapikey"} or "credential" in normalized

utils/serverless_config.py

@@ -1,451 +0,0 @@
-"""Validation and private storage for Runpod serverless worker settings.
-
-Secret storage is supported only on POSIX masters. The module remains
-importable elsewhere so ordinary worker configuration and redaction continue
-to work, but secret-store operations fail explicitly without POSIX locking.
-"""
-
-from __future__ import annotations
-
-import copy
-import json
-import os
-import re
-import secrets
-import stat
-import threading
-from contextlib import contextmanager
-from pathlib import Path
-from typing import Iterator
-
-from .secret_fields import is_secret_field_name
-
-try:
-    import fcntl
-except ImportError:  # Windows and other platforms without POSIX file locking.
-    fcntl = None
-
-
-class ServerlessConfigError(ValueError):
-    """A serverless worker configuration is invalid."""
-
-
-class SecretStoreError(RuntimeError):
-    """The secret store could not be accessed safely."""
-
-
-_ENDPOINT_ID_RE = re.compile(r"[A-Za-z0-9_-]{1,200}", re.ASCII)
-_UNSUPPORTED_SECRET_STORE_MESSAGE = (
-    "Serverless credential storage is currently supported on POSIX masters only"
-)
-# Images per run is also this worker's effective parallelism: every image job
-# is submitted at once (the runtime's fixed PER_WORKER_JOB_LIMIT equals this),
-# and the Runpod endpoint's max workers decides how many actually run together.
-MAX_JOBS_PER_RUN = 16
-_NUMERIC_FIELDS = {
-    "jobs_per_run": (1, MAX_JOBS_PER_RUN, 1),
-    "timeout_seconds": (30, 86400, 1800),
-}
-# Older configs saved a separate concurrency setting. It is accepted so they
-# still load, then dropped: PER_WORKER_JOB_LIMIT replaces it.
-_LEGACY_IGNORED_FIELDS = ("concurrency",)
-_THREAD_LOCKS: dict[str, threading.Lock] = {}
-_THREAD_LOCKS_GUARD = threading.Lock()
-
-
-def validate_serverless_worker(worker: dict) -> dict:
-    """Return a validated copy of a Runpod serverless worker record.
-
-    Provider credentials are intentionally not required here; they are stored
-    separately by :class:`SecretStore`.
-    """
-    if not isinstance(worker, dict):
-        raise ServerlessConfigError("worker must be an object")
-
-    try:
-        result = copy.deepcopy(worker)
-    except Exception:
-        raise ServerlessConfigError("worker must be copyable") from None
-
-    worker_id = result.get("id")
-    if not isinstance(worker_id, str) or not worker_id.strip():
-        raise ServerlessConfigError("id must be a non-empty string")
-
-    endpoint_id = result.get("endpoint_id")
-    if not isinstance(endpoint_id, str) or not _ENDPOINT_ID_RE.fullmatch(endpoint_id):
-        raise ServerlessConfigError(
-            "Endpoint ID must be just the Runpod endpoint ID (letters, numbers, - or _), not a URL"
-        )
-
-    if result.get("type", "serverless") != "serverless":
-        raise ServerlessConfigError("type must be serverless")
-    if result.get("provider", "runpod") != "runpod":
-        raise ServerlessConfigError("provider is not supported")
-    if "enabled" in result and type(result["enabled"]) is not bool:
-        raise ServerlessConfigError("enabled must be a boolean")
-
-    for field, (minimum, maximum, default) in _NUMERIC_FIELDS.items():
-        value = result.get(field, default)
-        if type(value) is not int or not minimum <= value <= maximum:
-            raise ServerlessConfigError(f"{field} must be an integer from {minimum} to {maximum}")
-        result[field] = value
-    for field in _LEGACY_IGNORED_FIELDS:
-        result.pop(field, None)
-
-    result["type"] = "serverless"
-    result["provider"] = "runpod"
-    return result
-
-
-def _thread_lock_for(path: str) -> threading.Lock:
-    with _THREAD_LOCKS_GUARD:
-        lock = _THREAD_LOCKS.get(path)
-        if lock is None:
-            lock = threading.Lock()
-            _THREAD_LOCKS[path] = lock
-        return lock
-
-
-class SecretStore:
-    """Atomically persist worker-id to API-key mappings in a private JSON file."""
-
-    def __init__(self, path: Path):
-        self.path = Path(path)
-
-    def get(self, worker_id: str) -> str | None:
-        self._require_supported()
-        worker_id = self._validate_worker_id(worker_id)
-        try:
-            with self._locked(create_parent=False) as (parent_fd, filename):
-                data = self._read_mapping(parent_fd, filename)
-                return data.get(worker_id)
-        except FileNotFoundError:
-            return None
-        except SecretStoreError:
-            raise
-        except Exception:
-            raise SecretStoreError("Unable to access secret store") from None
-
-    def set(self, worker_id: str, key: str) -> None:
-        self._require_supported()
-        worker_id = self._validate_worker_id(worker_id)
-        if (
-            not isinstance(key, str)
-            or not key
-            or len(key) > 4096
-            or "\r" in key
-            or "\n" in key
-        ):
-            raise SecretStoreError("secret must be 1 to 4096 characters and contain no CR or LF")
-        try:
-            with self._locked(create_parent=True) as (parent_fd, filename):
-                data = self._read_mapping(parent_fd, filename)
-                data[worker_id] = key
-                self._write_mapping(parent_fd, filename, data)
-        except SecretStoreError:
-            raise
-        except Exception:
-            raise SecretStoreError("Unable to access secret store") from None
-
-    def delete(self, worker_id: str) -> None:
-        self._require_supported()
-        worker_id = self._validate_worker_id(worker_id)
-        try:
-            with self._locked(create_parent=False) as (parent_fd, filename):
-                data = self._read_mapping(parent_fd, filename)
-                if worker_id in data:
-                    del data[worker_id]
-                    self._write_mapping(parent_fd, filename, data)
-        except FileNotFoundError:
-            return
-        except SecretStoreError:
-            raise
-        except Exception:
-            raise SecretStoreError("Unable to access secret store") from None
-
-    def has(self, worker_id: str) -> bool:
-        return self.get(worker_id) is not None
-
-    @staticmethod
-    def _require_supported() -> None:
-        # Never substitute a weaker Windows store: all operations retain the
-        # dir_fd/O_NOFOLLOW and flock protections implemented below.
-        if os.name != "posix" or fcntl is None:
-            raise SecretStoreError(_UNSUPPORTED_SECRET_STORE_MESSAGE)
-
-    @staticmethod
-    def _validate_worker_id(worker_id: str) -> str:
-        if not isinstance(worker_id, str) or not worker_id.strip():
-            raise SecretStoreError("worker_id must be a non-empty string")
-        return worker_id
-
-    def _absolute_path(self) -> Path:
-        path = self.path
-        if ".." in path.parts:
-            raise SecretStoreError("unsafe secret store path")
-        if not path.is_absolute():
-            path = Path.cwd() / path
-        if len(path.parts) < 2 or path.name in ("", ".", ".."):
-            raise SecretStoreError("unsafe secret store path")
-        return path
-
-    @contextmanager
-    def _locked(self, create_parent: bool) -> Iterator[tuple[int, str]]:
-        lock_api = fcntl
-        if os.name != "posix" or lock_api is None:
-            raise SecretStoreError(_UNSUPPORTED_SECRET_STORE_MESSAGE)
-        path = self._absolute_path()
-        filename = path.name
-        lock_name = f".{filename}.lock"
-        thread_lock = _thread_lock_for(str(path))
-
-        with thread_lock:
-            parent_fd = self._open_parent(path, create=create_parent)
-            lock_fd = None
-            try:
-                flags = os.O_RDWR | os.O_CREAT | os.O_CLOEXEC | os.O_NOFOLLOW
-                lock_fd = os.open(lock_name, flags, 0o600, dir_fd=parent_fd)
-                lock_stat = os.fstat(lock_fd)
-                if not stat.S_ISREG(lock_stat.st_mode):
-                    raise SecretStoreError("unsafe secret store lock")
-                os.fchmod(lock_fd, 0o600)
-                lock_api.flock(lock_fd, lock_api.LOCK_EX)
-                self._check_target(parent_fd, filename)
-                yield parent_fd, filename
-            except SecretStoreError:
-                raise
-            except Exception:
-                raise SecretStoreError("Unable to access secret store") from None
-            finally:
-                if lock_fd is not None:
-                    try:
-                        lock_api.flock(lock_fd, lock_api.LOCK_UN)
-                    except OSError:
-                        pass
-                    os.close(lock_fd)
-                os.close(parent_fd)
-
-    @staticmethod
-    def _open_parent(path: Path, create: bool) -> int:
-        parts = path.parts
-        current_fd = os.open(path.anchor, os.O_RDONLY | os.O_DIRECTORY | os.O_CLOEXEC)
-        try:
-            for part in parts[1:-1]:
-                created = False
-                try:
-                    next_fd = os.open(
-                        part,
-                        os.O_RDONLY | os.O_DIRECTORY | os.O_CLOEXEC | os.O_NOFOLLOW,
-                        dir_fd=current_fd,
-                    )
-                except FileNotFoundError:
-                    if not create:
-                        raise
-                    try:
-                        os.mkdir(part, 0o700, dir_fd=current_fd)
-                        created = True
-                    except FileExistsError:
-                        pass
-                    next_fd = os.open(
-                        part,
-                        os.O_RDONLY | os.O_DIRECTORY | os.O_CLOEXEC | os.O_NOFOLLOW,
-                        dir_fd=current_fd,
-                    )
-                if created:
-                    os.fchmod(next_fd, 0o700)
-                os.close(current_fd)
-                current_fd = next_fd
-            return current_fd
-        except Exception:
-            os.close(current_fd)
-            raise
-
-    @staticmethod
-    def _check_target(parent_fd: int, filename: str) -> None:
-        try:
-            target_stat = os.stat(filename, dir_fd=parent_fd, follow_symlinks=False)
-        except FileNotFoundError:
-            return
-        if not stat.S_ISREG(target_stat.st_mode):
-            raise SecretStoreError("unsafe secret store file")
-
-    @classmethod
-    def _read_mapping(cls, parent_fd: int, filename: str) -> dict[str, str]:
-        cls._check_target(parent_fd, filename)
-        try:
-            fd = os.open(
-                filename,
-                os.O_RDONLY | os.O_CLOEXEC | os.O_NOFOLLOW,
-                dir_fd=parent_fd,
-            )
-        except FileNotFoundError:
-            return {}
-        except Exception:
-            raise SecretStoreError("Unable to access secret store") from None
-
-        try:
-            file_stat = os.fstat(fd)
-            if not stat.S_ISREG(file_stat.st_mode):
-                raise SecretStoreError("unsafe secret store file")
-            os.fchmod(fd, 0o600)
-            with os.fdopen(fd, "rb") as stream:
-                raw = stream.read()
-            try:
-                mapping = json.loads(raw.decode("utf-8"), object_pairs_hook=cls._unique_object)
-            except Exception:
-                raise SecretStoreError("secret store is corrupt") from None
-            if not isinstance(mapping, dict) or any(
-                not isinstance(worker_id, str)
-                or not worker_id.strip()
-                or not isinstance(key, str)
-                or not key
-                for worker_id, key in mapping.items()
-            ):
-                raise SecretStoreError("secret store is corrupt")
-            return mapping
-        except SecretStoreError:
-            raise
-        except Exception:
-            raise SecretStoreError("Unable to access secret store") from None
-        finally:
-            try:
-                os.close(fd)
-            except OSError:
-                pass
-
-    @staticmethod
-    def _unique_object(pairs: list[tuple[str, object]]) -> dict[str, object]:
-        result: dict[str, object] = {}
-        for key, value in pairs:
-            if key in result:
-                raise ValueError("duplicate key")
-            result[key] = value
-        return result
-
-    @staticmethod
-    def _write_mapping(parent_fd: int, filename: str, mapping: dict[str, str]) -> None:
-        payload = json.dumps(mapping, ensure_ascii=False, sort_keys=True, separators=(",", ":")).encode("utf-8")
-        temp_name = f".{filename}.{secrets.token_hex(12)}.tmp"
-        temp_fd = None
-        try:
-            flags = os.O_WRONLY | os.O_CREAT | os.O_EXCL | os.O_CLOEXEC | os.O_NOFOLLOW
-            temp_fd = os.open(temp_name, flags, 0o600, dir_fd=parent_fd)
-            os.fchmod(temp_fd, 0o600)
-            with os.fdopen(temp_fd, "wb") as stream:
-                temp_fd = None
-                stream.write(payload)
-                stream.flush()
-                os.fsync(stream.fileno())
-            SecretStore._check_target(parent_fd, filename)
-            os.replace(temp_name, filename, src_dir_fd=parent_fd, dst_dir_fd=parent_fd)
-            os.fsync(parent_fd)
-        except SecretStoreError:
-            raise
-        except Exception:
-            raise SecretStoreError("Unable to access secret store") from None
-        finally:
-            if temp_fd is not None:
-                try:
-                    os.close(temp_fd)
-                except OSError:
-                    pass
-            try:
-                os.unlink(temp_name, dir_fd=parent_fd)
-            except OSError:
-                pass
-
-
-class ServerlessSecretChanges:
-    """Apply serverless secret changes and restore them unless committed.
-
-    Use this around the associated config transaction so a failed config save
-    compensates any secret-store writes or deletions made along the way.
-    """
-
-    def __init__(self):
-        self._originals: dict[tuple[int, str], tuple[SecretStore, str, str | None]] = {}
-        self._mutations: list[tuple[int, str]] = []
-        self._committed = False
-
-    def __enter__(self) -> "ServerlessSecretChanges":
-        return self
-
-    def __exit__(self, _exc_type, _exc_value, _traceback) -> bool:
-        if not self._committed:
-            self.rollback()
-        return False
-
-    def capture(self, store: SecretStore, worker_id: str) -> str | None:
-        """Capture and return the original key once per store/worker pair."""
-        mutation_id = (id(store), worker_id)
-        if mutation_id not in self._originals:
-            self._originals[mutation_id] = (store, worker_id, store.get(worker_id))
-        return self._originals[mutation_id][2]
-
-    def set(self, store: SecretStore, worker_id: str, key: str) -> None:
-        self._prepare_mutation(store, worker_id)
-        store.set(worker_id, key)
-
-    def delete(self, store: SecretStore, worker_id: str) -> None:
-        self._prepare_mutation(store, worker_id)
-        store.delete(worker_id)
-
-    def commit(self) -> None:
-        """Keep the secret changes after the corresponding config save."""
-        self._committed = True
-
-    def rollback(self) -> None:
-        """Best-effort restoration without masking the original route error."""
-        for mutation_id in reversed(self._mutations):
-            store, worker_id, prior_secret = self._originals[mutation_id]
-            try:
-                if prior_secret is None:
-                    store.delete(worker_id)
-                else:
-                    store.set(worker_id, prior_secret)
-            except Exception:
-                pass
-        self._mutations.clear()
-
-    def _prepare_mutation(self, store: SecretStore, worker_id: str) -> None:
-        self.capture(store, worker_id)
-        mutation_id = (id(store), worker_id)
-        if mutation_id not in self._mutations:
-            # Record before applying: a store operation may fail after replacing
-            # its file, so the exception path must still attempt compensation.
-            self._mutations.append(mutation_id)
-
-
-def _redact_credentials(value):
-    if isinstance(value, dict):
-        redacted = {}
-        for key, item in value.items():
-            if is_secret_field_name(key):
-                continue
-            redacted[key] = _redact_credentials(item)
-        return redacted
-    if isinstance(value, list):
-        return [_redact_credentials(item) for item in value]
-    if isinstance(value, tuple):
-        return tuple(_redact_credentials(item) for item in value)
-    return value
-
-
-def public_config(config, store: SecretStore):
-    """Return a deep-copied public configuration with credentials removed."""
-    result = _redact_credentials(copy.deepcopy(config))
-    if not isinstance(result, dict):
-        return result
-
-    workers = result.get("workers")
-    if isinstance(workers, list):
-        for worker in workers:
-            if isinstance(worker, dict) and worker.get("type") == "serverless":
-                worker_id = worker.get("id")
-                configured = (
-                    store.has(worker_id)
-                    if isinstance(worker_id, str) and worker_id.strip()
-                    else False
-                )
-                worker["api_key_configured"] = configured
-    return result

utils/transfer_codec.py

@@ -1,411 +0,0 @@
-"""
-Compressed worker -> master image transfer.
-
-Workers encode their IMAGE batch with one of the formats below and POST it to
-the master as multipart/form-data chunks; the master decodes each chunk back
-into per-frame tensors. The goal is to cut network egress (e.g. cloud workers
-returning results to a local master), so formats range from bit-exact
-lossless to strongly compressed video.
-
-Formats
--------
-legacy_png  Original behaviour: one JSON request per frame, uncompressed PNG
-            in base64. Kept for compatibility with unmodified masters.
-png         Lossless PNG per frame, sent as binary.
-ffv1        Lossless FFV1 video for the whole batch (good for frame sequences).
-jpeg, webp  Lossy per-frame stills.
-h264, h265, av1
-            Lossy video for the whole batch. Smallest by far for video
-            frame batches, where consecutive frames are similar.
-
-All formats quantize to 8-bit exactly like the legacy path did, so the
-lossless formats reproduce the legacy master tensors bit-for-bit.
-"""
-import io
-import math
-import os
-
-import numpy as np
-from PIL import Image
-
-LEGACY_FORMAT = "legacy_png"
-IMAGE_FORMATS = ("png", "jpeg", "webp")
-VIDEO_FORMATS = ("ffv1", "h264", "h265", "av1")
-LOSSLESS_FORMATS = (LEGACY_FORMAT, "png", "ffv1")
-TRANSFER_FORMATS = (LEGACY_FORMAT, "png", "ffv1", "jpeg", "webp", "h264", "h265", "av1")
-
-DEFAULT_TRANSFER_FORMAT = LEGACY_FORMAT
-DEFAULT_TRANSFER_QUALITY = 90
-
-TRANSFER_FORMAT_TOOLTIP = (
-    "How workers send images back to the master. "
-    "legacy_png: original uncompressed PNG (largest). "
-    "png / ffv1: lossless, identical result. "
-    "jpeg / webp: lossy stills. "
-    "h264 / h265 / av1: lossy video of the whole batch, smallest for frame sequences."
-)
-TRANSFER_QUALITY_TOOLTIP = (
-    "Quality for lossy transfer formats (ignored by lossless ones). "
-    "JPEG/WebP use it directly; video codecs map it to CRF "
-    "(90 -> h264/h265 CRF 13, av1 CRF 16)."
-)
-
-# encoder name, max CRF, container-friendly pixel format
-_VIDEO_CODECS = {
-    "ffv1": ("ffv1", None, "bgr0"),
-    "h264": ("libx264", 51, "yuv420p"),
-    "h265": ("libx265", 51, "yuv420p"),
-    "av1": ("libsvtav1", 63, "yuv420p"),
-}
-_VIDEO_RATE = 24
-# Upper bound on frames per video chunk, so one request never holds an
-# unbounded batch even when frames compress extremely well.
-MAX_VIDEO_CHUNK_FRAMES = 240
-# When sizing the next video chunk from observed bytes/frame, aim this far
-# below the limit so ordinary variation does not force a re-encode.
-VIDEO_CHUNK_SIZE_HEADROOM = 0.85
-
-
-class TransferCodecError(ValueError):
-    """Raised when a transfer payload cannot be encoded or decoded."""
-
-
-class FrameTooLargeError(TransferCodecError):
-    """A single encoded frame does not fit in one request.
-
-    Raised by encode_chunks before anything has been yielded for that frame;
-    callers that materialize all chunks first can therefore fail without
-    having delivered a partial batch.
-    """
-
-
-def normalize_format(value):
-    fmt = str(value or DEFAULT_TRANSFER_FORMAT).strip().lower()
-    if fmt not in TRANSFER_FORMATS:
-        raise TransferCodecError(
-            f"Unknown transfer format '{value}'. Expected one of: {', '.join(TRANSFER_FORMATS)}"
-        )
-    return fmt
-
-
-def normalize_quality(value):
-    try:
-        quality = int(value)
-    except (TypeError, ValueError):
-        quality = DEFAULT_TRANSFER_QUALITY
-    return max(1, min(100, quality))
-
-
-def quality_to_crf(quality, max_crf):
-    """Map 1-100 quality onto a codec CRF scale (0 = lossless/best).
-
-    The curve is chosen so the same quality number looks roughly the same
-    across JPEG/WebP and video codecs: 90 -> CRF 13 (h264/h265), 16 (av1);
-    80 -> 19 / 24; 100 -> 0.
-    """
-    quality = normalize_quality(quality)
-    return int(round(max_crf * (1.0 - quality / 100.0) ** 0.6))
-
-
-def tensor_batch_to_uint8(images):
-    """IMAGE tensor [B,H,W,C] float 0..1 -> uint8 numpy [B,H,W,3].
-
-    Uses the same truncating quantization as utils.image.tensor_to_pil so that
-    lossless formats reproduce the legacy transfer exactly, but clips first so
-    out-of-range values no longer wrap around.
-    """
-    arr = images.detach().cpu().float().numpy()
-    if arr.ndim == 3:
-        arr = arr[None, ...]
-    if arr.ndim != 4:
-        raise TransferCodecError(f"Expected IMAGE tensor [B,H,W,C], got shape {tuple(arr.shape)}")
-    channels = arr.shape[-1]
-    if channels == 1:
-        arr = np.repeat(arr, 3, axis=-1)
-    elif channels >= 3:
-        arr = arr[..., :3]
-    else:
-        raise TransferCodecError(f"Unsupported channel count {channels}")
-    return np.ascontiguousarray(np.clip(255.0 * arr, 0, 255).astype(np.uint8))
-
-
-def uint8_frame_to_tensor_array(frame):
-    """uint8 [H,W,3] -> float32 [1,H,W,3] numpy, matching utils.image.pil_to_tensor."""
-    return (frame.astype(np.float32) / 255.0)[None, ...]
-
-
-# --------------------------------------------------------------------------- #
-# Still images
-# --------------------------------------------------------------------------- #
-
-def encode_image_frame(frame, fmt, quality):
-    img = Image.fromarray(frame)
-    quality = normalize_quality(quality)
-    buf = io.BytesIO()
-    if fmt == "png":
-        img.save(buf, format="PNG", compress_level=4)  # ~same size as 6, 2x faster
-    elif fmt == "jpeg":
-        img.save(buf, format="JPEG", quality=quality, subsampling=0 if quality >= 90 else 2)
-    elif fmt == "webp":
-        img.save(buf, format="WEBP", quality=quality, method=4)
-    else:
-        raise TransferCodecError(f"'{fmt}' is not a still-image transfer format")
-    return buf.getvalue()
-
-
-def decode_image_frame(data):
-    try:
-        with Image.open(io.BytesIO(data)) as img:
-            return np.asarray(img.convert("RGB"), dtype=np.uint8)
-    except Exception as exc:
-        raise TransferCodecError(f"Failed to decode image frame: {exc}") from exc
-
-
-# --------------------------------------------------------------------------- #
-# Video
-# --------------------------------------------------------------------------- #
-
-def _import_av():
-    # SVT-AV1 prints a banner per encoder instance unless told otherwise.
-    os.environ.setdefault("SVT_LOG", "1")
-    try:
-        import av  # PyAV ships with ComfyUI
-    except ImportError as exc:
-        raise TransferCodecError("PyAV ('av') is required for video transfer formats") from exc
-    return av
-
-
-def video_codec_available(fmt):
-    if fmt not in _VIDEO_CODECS:
-        return False
-    try:
-        av = _import_av()
-        av.codec.Codec(_VIDEO_CODECS[fmt][0], "w")
-        return True
-    except Exception:
-        return False
-
-
-def _padded_size(fmt, pix_fmt, width, height):
-    """Encoder-friendly frame size; the decoder crops back to width x height.
-
-    4:2:0 needs even dimensions, and libx265 crashes on frames smaller than
-    its 64px CTU, so pad up to whichever is larger.
-    """
-    if pix_fmt == "yuv420p":
-        width, height = width + width % 2, height + height % 2
-    if fmt == "h265":
-        width, height = max(width, 64), max(height, 64)
-    return width, height
-
-
-def _pad_to(frame, width, height):
-    h, w = frame.shape[:2]
-    if h == height and w == width:
-        return frame
-    return np.pad(frame, ((0, height - h), (0, width - w), (0, 0)), mode="edge")
-
-
-class _VideoChunkWriter:
-    def __init__(self, fmt, quality, width, height):
-        av = _import_av()
-        self._av = av
-        codec_name, max_crf, pix_fmt = _VIDEO_CODECS[fmt]
-        self.fmt = fmt
-        self.pix_fmt = pix_fmt
-        self.width, self.height = _padded_size(fmt, pix_fmt, width, height)
-        self.buffer = io.BytesIO()
-        self.container = av.open(self.buffer, mode="w", format="matroska")
-        self.stream = self.container.add_stream(codec_name, rate=_VIDEO_RATE)
-        self.stream.width = self.width
-        self.stream.height = self.height
-        self.stream.pix_fmt = pix_fmt
-        options = {}
-        if max_crf is not None:
-            options["crf"] = str(quality_to_crf(quality, max_crf))
-        if fmt == "h264":
-            options["preset"] = "medium"
-        elif fmt == "h265":
-            options["preset"] = "fast"
-            options["x265-params"] = "log-level=error"
-        elif fmt == "av1":
-            options["preset"] = "8"
-        elif fmt == "ffv1":
-            options["level"] = "3"
-            options["slicecrc"] = "0"
-        self.stream.options = options
-        self.count = 0
-
-    def add(self, frame):
-        frame = _pad_to(frame, self.width, self.height)
-        vf = self._av.VideoFrame.from_ndarray(frame, format="rgb24")
-        for packet in self.stream.encode(vf):
-            self.container.mux(packet)
-        self.count += 1
-
-    def bytes_so_far(self):
-        return self.buffer.tell()
-
-    def close(self):
-        for packet in self.stream.encode():
-            self.container.mux(packet)
-        self.container.close()
-        return self.buffer.getvalue()
-
-
-def decode_video(data, width, height, expected_frames=None):
-    av = _import_av()
-    frames = []
-    try:
-        with av.open(io.BytesIO(data), mode="r") as container:
-            for vf in container.decode(video=0):
-                arr = vf.to_ndarray(format="rgb24")
-                frames.append(np.ascontiguousarray(arr[:height, :width]))
-    except TransferCodecError:
-        raise
-    except Exception as exc:
-        raise TransferCodecError(f"Failed to decode video chunk: {exc}") from exc
-    if expected_frames is not None and len(frames) != expected_frames:
-        raise TransferCodecError(
-            f"Video chunk decoded to {len(frames)} frames, expected {expected_frames}"
-        )
-    return frames
-
-
-# --------------------------------------------------------------------------- #
-# Chunking
-# --------------------------------------------------------------------------- #
-
-def encode_chunks(frames, fmt, quality, max_chunk_bytes, max_frame_bytes=None):
-    """Yield transfer chunks for a uint8 frame batch [B,H,W,3].
-
-    Each chunk is a dict:
-        {"format", "start_index", "count", "width", "height", "parts": [bytes, ...]}
-    Image formats carry one part per frame; video formats carry one part
-    (a Matroska file) holding `count` frames.
-
-    Size guarantees, measured on the finalized bytes (after encoder flush and
-    muxer finalization, not on encoder buffers):
-
-    * a chunk holding more than one frame never exceeds `max_chunk_bytes`;
-    * a chunk holding a single frame may exceed `max_chunk_bytes` but never
-      `max_frame_bytes` (defaults to `max_chunk_bytes`).
-
-    A frame that cannot fit in `max_frame_bytes` even on its own raises
-    FrameTooLargeError. Video chunks that come out too large are re-encoded
-    with fewer frames.
-    """
-    fmt = normalize_format(fmt)
-    if fmt == LEGACY_FORMAT:
-        raise TransferCodecError("legacy_png is sent through the JSON path, not encode_chunks")
-    total = int(frames.shape[0])
-    if total == 0:
-        return
-    height, width = int(frames.shape[1]), int(frames.shape[2])
-    max_chunk_bytes = max(1, int(max_chunk_bytes))
-    max_frame_bytes = max_chunk_bytes if max_frame_bytes is None else max(max_chunk_bytes, int(max_frame_bytes))
-
-    def _chunk(start, parts, count):
-        return {
-            "format": fmt,
-            "start_index": start,
-            "count": count,
-            "width": width,
-            "height": height,
-            "parts": parts,
-        }
-
-    def _too_large(index, size):
-        return FrameTooLargeError(
-            f"Frame {index} encodes to {size} bytes as {fmt}, above the {max_frame_bytes}-byte request limit. "
-            "Use a lossy format or lower quality, or raise the master's --max-upload-size together with "
-            "COMFYUI_DISTRIBUTED_TRANSFER_MAX_REQUEST_BYTES."
-        )
-
-    if fmt in IMAGE_FORMATS:
-        start, parts, size = 0, [], 0
-        for i in range(total):
-            data = encode_image_frame(frames[i], fmt, quality)
-            if len(data) > max_frame_bytes:
-                raise _too_large(i, len(data))
-            if parts and size + len(data) > max_chunk_bytes:
-                yield _chunk(start, parts, len(parts))
-                start, parts, size = i, [], 0
-            parts.append(data)
-            size += len(data)
-        if parts:
-            yield _chunk(start, parts, len(parts))
-        return
-
-    start = 0
-    budget = MAX_VIDEO_CHUNK_FRAMES
-    while start < total:
-        want = min(budget, total - start)
-        data, count = _encode_video_range(frames, start, want, fmt, quality, width, height, max_chunk_bytes)
-        if len(data) <= max_chunk_bytes or (count == 1 and len(data) <= max_frame_bytes):
-            yield _chunk(start, [data], count)
-            start += count
-            per_frame = len(data) / count
-            budget = max(1, min(MAX_VIDEO_CHUNK_FRAMES, int(max_chunk_bytes * VIDEO_CHUNK_SIZE_HEADROOM / per_frame)))
-            continue
-        if count == 1:
-            raise _too_large(start, len(data))
-        # Finalized chunk overshot: retry this range with proportionally fewer
-        # frames. Strictly fewer than `count`, so this always terminates.
-        estimate = int(count * max_chunk_bytes * VIDEO_CHUNK_SIZE_HEADROOM / len(data))
-        budget = max(1, min(count - 1, estimate))
-
-
-def _encode_video_range(frames, start, count, fmt, quality, width, height, stop_after_bytes):
-    """Encode frames[start:start+count] as one finalized video.
-
-    Stops feeding frames early once the muxed output already exceeds
-    `stop_after_bytes`, since the chunk is then known to be too large and will
-    be re-encoded smaller; this bounds wasted work. Encoder lookahead means the
-    pre-flush size underestimates the final size, so the caller must still
-    check the returned (finalized) bytes. Returns (bytes, frames_encoded).
-    """
-    writer = _VideoChunkWriter(fmt, quality, width, height)
-    for i in range(start, start + count):
-        writer.add(frames[i])
-        if writer.bytes_so_far() > stop_after_bytes:
-            break
-    return writer.close(), writer.count
-
-
-def decode_chunk(envelope, parts):
-    """Decode one received chunk into a list of uint8 [H,W,3] frames."""
-    fmt = normalize_format(envelope.get("format"))
-    if fmt == LEGACY_FORMAT:
-        raise TransferCodecError("legacy_png chunks are not sent as multipart")
-    count = envelope.get("count")
-    width = envelope.get("width")
-    height = envelope.get("height")
-    for name, value in (("count", count), ("width", width), ("height", height)):
-        if not isinstance(value, int) or isinstance(value, bool) or value <= 0:
-            raise TransferCodecError(f"envelope.{name}: expected positive integer")
-    if not parts:
-        raise TransferCodecError("chunk has no media parts")
-
-    if fmt in IMAGE_FORMATS:
-        if len(parts) != count:
-            raise TransferCodecError(f"expected {count} image parts, got {len(parts)}")
-        frames = [decode_image_frame(p) for p in parts]
-    else:
-        if len(parts) != 1:
-            raise TransferCodecError(f"expected 1 video part, got {len(parts)}")
-        frames = decode_video(parts[0], width, height, expected_frames=count)
-
-    for frame in frames:
-        if frame.shape[0] != height or frame.shape[1] != width:
-            raise TransferCodecError(
-                f"decoded frame is {frame.shape[1]}x{frame.shape[0]}, expected {width}x{height}"
-            )
-    return frames
-
-
-def describe_ratio(sent_bytes, raw_bytes):
-    if sent_bytes <= 0:
-        return "n/a"
-    ratio = raw_bytes / sent_bytes
-    return f"{ratio:.1f}x" if math.isfinite(ratio) else "n/a"

web/apiClient.js

@@ -7,7 +7,7 @@ export function createApiClient(baseUrl) {
     const request = async (
         endpoint,
         options = {},
-        { retries = TIMEOUTS.MAX_RETRIES, retry = true, redactErrors = false } = {},
+        { retries = TIMEOUTS.MAX_RETRIES, retry = true } = {},
     ) => {
         const maxAttempts = retry ? retries : 1;
         let lastError;
@@ -36,7 +36,7 @@ export function createApiClient(baseUrl) {
                 return await response.json();
             } catch (error) {
                 lastError = error;
-                console.log(`API Error (attempt ${attempt + 1}/${maxAttempts}): ${endpoint} - ${redactErrors ? "[redacted]" : error.message}`);
+                console.log(`API Error (attempt ${attempt + 1}/${maxAttempts}): ${endpoint} - ${error.message}`);
                 if (attempt < maxAttempts - 1) {
                     await new Promise(resolve => setTimeout(resolve, delay));
                     delay *= 2; // Exponential backoff
@@ -86,35 +86,9 @@ export function createApiClient(baseUrl) {
         },
         
         async updateWorker(workerId, data) {
-            const apiKey = data?.api_key;
-            try {
-                return await request('/distributed/config/update_worker', {
-                    method: 'POST',
-                    body: JSON.stringify({ worker_id: workerId, ...data })
-                }, { retry: false, redactErrors: Boolean(apiKey) });
-            } catch (error) {
-                if (apiKey && typeof error?.message === "string") {
-                    error.message = error.message.split(apiKey).join("[redacted]");
-                }
-                throw error;
-            }
-        },
-
-        async testServerlessWorker(workerId) {
-            return request('/distributed/serverless/test', {
+            return request('/distributed/config/update_worker', {
                 method: 'POST',
-                body: JSON.stringify({ worker_id: workerId }),
-            }, { retry: false });
-        },
-
-        async getServerlessStatus(workerId) {
-            return request(`/distributed/serverless/status/${encodeURIComponent(workerId)}`);
-        },
-
-        async cancelServerlessWorker(workerId) {
-            return request('/distributed/serverless/cancel', {
-                method: 'POST',
-                body: JSON.stringify({ worker_id: workerId }),
+                body: JSON.stringify({ worker_id: workerId, ...data })
             }, { retry: false });
         },
         

web/constants.js

@@ -104,8 +104,7 @@ export const PULSE_ANIMATION_CSS = `
         transition: max-height 0.3s ease, opacity 0.3s ease, padding 0.3s ease, margin 0.3s ease;
     }
     .worker-settings.expanded {
-        /* Let long forms and wrapped labels fit; the sidebar owns scrolling. */
-        max-height: none;
+        max-height: 500px;
         opacity: 1;
         padding: 12px;
         margin-top: 8px;
@@ -131,7 +130,7 @@ export const PULSE_ANIMATION_CSS = `
 `;
 
 export const UI_STYLES = {
-    statusDot: "display: inline-block; flex: 0 0 10px; width: 10px; height: 10px; border-radius: 50%; margin-right: 10px;",
+    statusDot: "display: inline-block; width: 10px; height: 10px; border-radius: 50%; margin-right: 10px;",
     controlsDiv: "padding: 0 12px 12px 12px; display: flex; gap: 6px;",
     formGroup: "display: flex; flex-direction: column; gap: 5px;",
     formLabel: "font-size: 12px; color: var(--dist-label-text, #ccc); font-weight: 500;",
@@ -147,12 +146,12 @@ export const UI_STYLES = {
     // Column styles
     columnBase: "display: flex; align-items: center; justify-content: center;",
     checkboxColumn: "flex: 0 0 44px; display: flex; align-items: center; justify-content: center; cursor: default;",
-    contentColumn: "flex: 1; min-width: 0; display: flex; flex-direction: column; transition: background-color 0.2s ease;",
+    contentColumn: "flex: 1; display: flex; flex-direction: column; transition: background-color 0.2s ease;",
     iconColumn: "width: 44px; flex-shrink: 0; font-size: 20px; color: var(--dist-placeholder-add-color, #666);",
     
     // Row and content styles
     infoRow: "display: flex; align-items: center; padding: 12px; cursor: pointer; min-height: 64px;",
-    workerContent: "display: flex; align-items: center; gap: 10px; flex: 1; min-width: 0;",
+    workerContent: "display: flex; align-items: center; gap: 10px; flex: 1;",
     
     // Form and controls styles
     buttonGroup: "display: flex; gap: 4px; margin-top: 10px;",
@@ -225,7 +224,6 @@ export const ENDPOINTS = {
 
 export const NODE_CLASSES = {
     DISTRIBUTED_COLLECTOR: 'DistributedCollector',
-    IMAGE_LIST_START: 'DistributedImageListStart',
     DISTRIBUTED_SEED: 'DistributedSeed',
     DISTRIBUTED_EMPTY_IMAGE: 'DistributedEmptyImage',
     UPSCALE_DISTRIBUTED: 'UltimateSDUpscaleDistributed',

web/distributed.css

@@ -235,17 +235,6 @@
     border-radius: 4px;
 }
 
-.dist-worker-info {
-    min-width: 0;
-    overflow-wrap: anywhere;
-}
-
-.dist-serverless-status {
-    display: block;
-    margin-top: 4px;
-    line-height: 1.4;
-}
-
 .dist-worker-info__title {
     color: var(--dist-card-title) !important;
     font-size: 1.03em;

web/distributedValue.js

@@ -2,6 +2,7 @@ import { app } from "/scripts/app.js";
 import { ENDPOINTS } from "./constants.js";
 
 const NODE_CLASS = "DistributedValue";
+const CONVERTED_WIDGET = "converted-widget";
 const DYNAMIC_DEFAULT_WIDGET = "_dv_default";
 const DYNAMIC_WORKER_WIDGET_PREFIX = "_dv_worker_";
 const WORKERS_CHANGED_EVENT = "distributed:workers-changed";
@@ -41,13 +42,15 @@ function getDynamicWorkerWidgets(node) {
     return (node.widgets || []).filter((w) => w.name.startsWith(DYNAMIC_WORKER_WIDGET_PREFIX));
 }
 
-function hideWidgetForGood(widget) {
+function hideWidgetForGood(node, widget, suffix = "") {
     if (!widget) return;
-    // These are serialized backing fields, not widgets converted into sockets.
-    // Keep their native type/serializer and mutate the shared options in place
-    // so both the classic canvas and Nodes 2.0 suppress the whole widget row.
-    widget.options ??= {};
-    widget.options.hidden = true;
+    if (typeof widget.type === "string" && widget.type.startsWith(CONVERTED_WIDGET)) return;
+
+    widget.origType = widget.type;
+    widget.origComputeSize = widget.computeSize;
+    widget.origSerializeValue = widget.serializeValue;
+    widget.computeSize = () => [0, -4];
+    widget.type = `${CONVERTED_WIDGET}${suffix}`;
 
     // Hide any attached DOM element (multiline widgets).
     if (widget.element) widget.element.style.display = "none";
@@ -55,14 +58,14 @@ function hideWidgetForGood(widget) {
 
     if (widget.linkedWidgets) {
         for (const linked of widget.linkedWidgets) {
-            hideWidgetForGood(linked);
+            hideWidgetForGood(node, linked, `:${widget.name}`);
         }
     }
 }
 
 function hideRawWidgets(node) {
-    hideWidgetForGood(getRawDefaultWidget(node));
-    hideWidgetForGood(getRawWorkerValuesWidget(node));
+    hideWidgetForGood(node, getRawDefaultWidget(node), ":default_value");
+    hideWidgetForGood(node, getRawWorkerValuesWidget(node), ":worker_values");
 }
 
 function removeDynamicDefaultWidget(node) {

web/executionUtils.js

@@ -9,9 +9,7 @@ export function setupInterceptor(extension) {
             const hasCollector = findNodesByClass(prompt.output, NODE_CLASSES.DISTRIBUTED_COLLECTOR).length > 0;
             const hasDistUpscale = findNodesByClass(prompt.output, NODE_CLASSES.UPSCALE_DISTRIBUTED).length > 0;
 
-            const hasImageList = findNodesByClass(prompt.output, NODE_CLASSES.IMAGE_LIST_START).length > 0;
-
-            if (hasCollector || hasDistUpscale || hasImageList) {
+            if (hasCollector || hasDistUpscale) {
                 const result = await executeParallelDistributed(extension, prompt);
                 // Immediate status check for instant feedback
                 checkAllWorkerStatuses(extension);
@@ -27,22 +25,14 @@ export function setupInterceptor(extension) {
 export async function executeParallelDistributed(extension, promptWrapper) {
     const traceExecutionId = `exec_${Date.now()}_${generateUUID().slice(0, 6)}`;
     try {
-        const hasImageList = findNodesByClass(promptWrapper.output, NODE_CLASSES.IMAGE_LIST_START).length > 0;
-        // Image-list v1 uses direct workers only; do not contact paid provider pools.
-        const enabledWorkers = hasImageList
-            ? extension.enabledWorkers.filter(worker => worker.type !== "serverless")
-            : extension.enabledWorkers;
+        const enabledWorkers = extension.enabledWorkers;
         extension.log(`[exec:${traceExecutionId}] Starting distributed execution`, "debug");
         
         // Pre-flight health check on all enabled workers
         const activeWorkers = await performPreflightCheck(extension, enabledWorkers);
         
-        if (hasImageList && activeWorkers.length === 0 && extension.config?.settings?.master_delegate_only) {
-            throw new Error("Distributed Image List in orchestrator-only mode requires an active direct worker.");
-        }
-
         // Case: Enabled workers but all offline
-        if (!hasImageList && activeWorkers.length === 0 && enabledWorkers.length > 0) {
+        if (activeWorkers.length === 0 && enabledWorkers.length > 0) {
             extension.log("No active workers found. All enabled workers are offline.");
             if (extension.ui?.showToast) {
                 extension.ui.showToast(extension.app, "error", "All Workers Offline", 
@@ -122,21 +112,6 @@ export async function performPreflightCheck(extension, workers) {
     const startTime = Date.now();
     
     const checkPromises = workers.map(async (worker) => {
-        if (worker.type === "serverless") {
-            try {
-                const providerStatus = await extension.api.getServerlessStatus(worker.id);
-                if (providerStatus?.ok) {
-                    extension.log(`Serverless worker ${worker.name} is ready to scale`, "debug");
-                    return { worker, active: true };
-                }
-                extension.log(`Serverless worker ${worker.name} is not ready: ${providerStatus?.message || "unavailable"}`, "debug");
-                return { worker, active: false };
-            } catch (error) {
-                extension.log(`Serverless worker ${worker.name} status check failed: ${error.message}`, "debug");
-                return { worker, active: false };
-            }
-        }
-
         const workerUrl = getWorkerUrl(extension, worker);
         
         extension.log(`Pre-flight checking ${worker.name} at: ${workerUrl}`, "debug");

web/imageListStart.js

@@ -1,170 +0,0 @@
-import { app } from "/scripts/app.js";
-
-// Distributed Image List Start declares value_1..value_10 outputs in Python so
-// saved links and prompt output indexes stay valid. ComfyUI grows the matching
-// `values.value_N` inputs natively (Autogrow) but has no growable outputs, so
-// this extension only presents them: it shows outputs for connected inputs,
-// copies each input's type onto its output, and moves output links when
-// Autogrow compacts inputs after a middle input is disconnected.
-// If this script fails, the node still runs; all outputs simply stay visible.
-
-const START_CLASS = "DistributedImageListStart";
-const FIXED_OUTPUTS = 4; // image, mask, item_index, context
-const MAX_VALUES = 10;
-const INPUT_PREFIX = "values.value_";
-
-function valueOrdinal(name) {
-    if (typeof name !== "string" || !name.startsWith(INPUT_PREFIX)) return null;
-    const ordinal = Number.parseInt(name.slice(INPUT_PREFIX.length), 10);
-    return ordinal >= 1 && ordinal <= MAX_VALUES ? ordinal : null;
-}
-
-function getLink(graph, id) {
-    if (id == null || !graph) return undefined;
-    if (typeof graph.getLink === "function") return graph.getLink(id) ?? undefined;
-    const links = graph.links;
-    return (typeof links?.get === "function" ? links.get(id) : links?.[id]) ?? undefined;
-}
-
-function connectedValues(node) {
-    const result = new Map();
-    for (const input of node.inputs ?? []) {
-        const ordinal = valueOrdinal(input.name);
-        if (!ordinal) continue;
-        const link = getLink(node.graph, input.link);
-        if (!link) continue;
-        const origin = node.graph.getNodeById?.(link.origin_id);
-        const originType = origin?.outputs?.[link.origin_slot]?.type;
-        const type = originType && originType !== "*" ? originType : (link.type ?? "*");
-        result.set(ordinal, { type, key: `${link.origin_id}:${link.origin_slot}` });
-    }
-    return result;
-}
-
-function outputLinks(node, ordinal) {
-    return node.outputs?.[FIXED_OUTPUTS + ordinal - 1]?.links ?? [];
-}
-
-function typesCompatible(outputType, inputType) {
-    if (!outputType || !inputType || outputType === "*" || inputType === "*") return true;
-    const accepted = String(inputType).split(",");
-    return String(outputType).split(",").some(type => accepted.includes(type));
-}
-
-function moveOutputLinks(node, from, to) {
-    const graph = node.graph;
-    const targets = outputLinks(node, from)
-        .map(id => getLink(graph, id))
-        .filter(Boolean)
-        .map(link => [graph.getNodeById?.(link.target_id), link.target_slot])
-        .filter(([target]) => target);
-    // The moved links carry the same value, so give the destination output its
-    // type first; LiteGraph refuses connections to a mismatched output type.
-    node.outputs[FIXED_OUTPUTS + to - 1].type = node.outputs[FIXED_OUTPUTS + from - 1].type;
-    // Connecting a target input replaces its existing link from `from`.
-    for (const [target, slot] of targets) node.connect(FIXED_OUTPUTS + to - 1, target, slot);
-}
-
-function disconnectValueOutput(node, ordinal) {
-    if (outputLinks(node, ordinal).length) node.disconnectOutput(FIXED_OUTPUTS + ordinal - 1);
-}
-
-// Autogrow removes a disconnected middle input and shifts later input links
-// down by one. Shift the matching outputs' links the same way, and drop the
-// removed value's downstream links, so no consumer silently receives a
-// different value.
-function followCompactedInputs(node, current) {
-    const previous = node._distributedValueSources;
-    if (!previous) return;
-    const now = [];
-    for (const [ordinal, info] of current) now[ordinal - 1] = info.key;
-    const before = previous;
-    if (now.length !== before.length - 1) return;
-    let removed = 0;
-    while (removed < now.length && now[removed] === before[removed]) removed++;
-    // Removing the last value shifts nothing; keep its links so a temporary
-    // rewire does not silently drop downstream connections.
-    if (removed >= now.length) return;
-    for (let i = removed; i < now.length; i++) {
-        if (now[i] !== before[i + 1]) return; // Not a single Autogrow compaction.
-    }
-    disconnectValueOutput(node, removed + 1);
-    for (let i = removed + 1; i < before.length; i++) moveOutputLinks(node, i + 1, i);
-}
-
-function syncImageListStartOutputs(node) {
-    if (!node?.graph || !Array.isArray(node.outputs) || node.outputs.length < FIXED_OUTPUTS) return;
-    const current = connectedValues(node);
-    followCompactedInputs(node, current);
-    const sources = [];
-    for (const [ordinal, info] of current) sources[ordinal - 1] = info.key;
-    node._distributedValueSources = Array.from(sources);
-
-    let visible = Math.max(0, ...current.keys());
-    for (let ordinal = MAX_VALUES; ordinal > visible; ordinal--) {
-        // Never drop a wired output; its links keep it visible until removed.
-        if (outputLinks(node, ordinal).length) visible = ordinal;
-    }
-    const listShape = node.outputs[0]?.shape;
-    while (node.outputs.length < FIXED_OUTPUTS + visible) {
-        const ordinal = node.outputs.length - FIXED_OUTPUTS + 1;
-        node.addOutput(`value_${ordinal}`, "*", listShape === undefined ? undefined : { shape: listShape });
-    }
-    while (node.outputs.length > FIXED_OUTPUTS + visible) node.removeOutput(node.outputs.length - 1);
-
-    for (let ordinal = 1; ordinal <= visible; ordinal++) {
-        const output = node.outputs[FIXED_OUTPUTS + ordinal - 1];
-        const type = current.get(ordinal)?.type ?? "*";
-        output.name = `value_${ordinal}`;
-        if (output.type !== type) {
-            output.type = type;
-            for (const id of [...(output.links ?? [])]) {
-                const link = getLink(node.graph, id);
-                if (!link) continue;
-                const target = node.graph.getNodeById?.(link.target_id);
-                const targetType = target?.inputs?.[link.target_slot]?.type;
-                if (target && !typesCompatible(type, targetType)) target.disconnectInput(link.target_slot);
-                else link.type = type;
-            }
-        }
-    }
-    node.setDirtyCanvas?.(true, true);
-}
-
-function scheduleSync(node) {
-    if (node._distributedValueSyncPending) return;
-    node._distributedValueSyncPending = true;
-    const run = () => {
-        node._distributedValueSyncPending = false;
-        try {
-            syncImageListStartOutputs(node);
-        } catch (error) {
-            console.error("[Distributed] Image List Start output sync failed:", error);
-        }
-    };
-    // Run after Autogrow's own animation-frame input compaction.
-    if (typeof requestAnimationFrame === "function") requestAnimationFrame(() => setTimeout(run, 0));
-    else setTimeout(run, 0);
-}
-
-app.registerExtension({
-    name: "Distributed.ImageListStartValues",
-    nodeCreated(node) {
-        if (node.comfyClass !== START_CLASS) return;
-        const onConnectionsChange = node.onConnectionsChange;
-        node.onConnectionsChange = function (...args) {
-            const result = onConnectionsChange?.apply(this, args);
-            scheduleSync(this);
-            return result;
-        };
-        const onConfigure = node.onConfigure;
-        node.onConfigure = function (...args) {
-            const result = onConfigure?.apply(this, args);
-            // A loaded graph is the source of truth for prior wiring.
-            this._distributedValueSources = undefined;
-            scheduleSync(this);
-            return result;
-        };
-        scheduleSync(node);
-    },
-});

web/main.js

@@ -15,8 +15,7 @@ import { parseHostInput, getMasterUrl as buildMasterUrl } from './urlUtils.js';
 
 const WORKERS_CHANGED_EVENT = "distributed:workers-changed";
 
-// Exported for tests; ComfyUI only uses the registerExtension side effect below.
-export class DistributedExtension {
+class DistributedExtension {
     constructor() {
         this.config = null;
         this.originalQueuePrompt = api.queuePrompt.bind(api);
@@ -197,7 +196,6 @@ export class DistributedExtension {
 
     async updateWorkerEnabled(workerId, enabled) {
         const worker = this.config.workers.find(w => w.id === workerId);
-        const previousEnabled = worker ? Boolean(worker.enabled) : !enabled;
         if (worker) {
             worker.enabled = enabled;
             this.state.updateWorker(workerId, { enabled });
@@ -221,19 +219,6 @@ export class DistributedExtension {
             await this.api.updateWorker(workerId, { enabled });
         } catch (error) {
             this.log("Error updating worker: " + error.message, "error");
-            // The backend kept the old state (e.g. a serverless worker without a
-            // saved API key); show that instead of a toggle that did not stick.
-            if (worker) {
-                worker.enabled = previousEnabled;
-                this.state.updateWorker(workerId, { enabled: previousEnabled });
-                this._emitWorkersChanged();
-            }
-            this.app?.extensionManager?.toast?.add({
-                severity: "error",
-                summary: enabled ? "Could not enable worker" : "Could not disable worker",
-                detail: error.message,
-                life: 5000,
-            });
         }
 
         if (this.panelElement) {

web/masterDetection.js

@@ -42,15 +42,6 @@ export async function detectMasterIP(extension) {
             if (shouldAutoPopulate) {
                 extension.log(`Auto-populating workers based on ${extension.cudaDeviceCount} CUDA devices (excluding master on CUDA ${extension.masterCudaDevice})`, "info");
 
-                const workerCount = extension.cudaDeviceCount - (
-                    Number.isInteger(extension.masterCudaDevice) &&
-                    extension.masterCudaDevice >= 0 &&
-                    extension.masterCudaDevice < extension.cudaDeviceCount ? 1 : 0
-                );
-                const workerPorts = data.local_worker_ports;
-                if (!Array.isArray(workerPorts) || workerPorts.length !== workerCount) {
-                    throw new Error("Could not allocate enough available local worker ports; check the server log");
-                }
                 const newWorkers = [];
                 let workerNum = 1;
                 let portOffset = 0;
@@ -65,7 +56,7 @@ export async function detectMasterIP(extension) {
                         id: generateUUID(),
                         name: `Worker ${workerNum}`,
                         host: isRunpod ? null : "localhost",
-                        port: workerPorts[portOffset],
+                        port: 8189 + portOffset,
                         cuda_device: i,
                         enabled: true,
                         extra_args: isRunpod ? "--listen" : "",

web/sidebarRenderer.js

@@ -18,19 +18,6 @@ export function updateWorkerCard(extension, workerId, newStatus = {}) {
     const workerState = extension.state.getWorker(workerId);
     const isLaunching = Boolean(workerState?.launching);
 
-    if (worker.type === "serverless") {
-        const providerStatus = workerState?.serverlessStatus || {};
-        if (newStatus.online && newStatus.processing) {
-            extension.ui.updateStatusDot(workerId, STATUS_COLORS.PROCESSING_YELLOW, `Runpod active (${newStatus.queueCount || 0} queued)`, false);
-        } else if (newStatus.online) {
-            extension.ui.updateStatusDot(workerId, STATUS_COLORS.ONLINE_GREEN, providerStatus.message || "Runpod ready to scale", false);
-        } else if (worker.enabled) {
-            extension.ui.updateStatusDot(workerId, STATUS_COLORS.OFFLINE_RED, providerStatus.message || "Runpod endpoint unavailable", false);
-        }
-        updateWorkerControls(extension, workerId);
-        return true;
-    }
-
     if (isLaunching && !newStatus.online) {
         extension.ui.updateStatusDot(workerId, STATUS_COLORS.PROCESSING_YELLOW, "Launching...", true);
     } else if (newStatus.online && newStatus.processing) {

web/tests/distributedValue.test.js

@@ -1,120 +0,0 @@
-import { readFileSync } from "node:fs";
-import vm from "node:vm";
-import { describe, expect, it, vi } from "vitest";
-
-// Evaluate this repository's browser entrypoint without needing a ComfyUI server.
-// The real-renderer acceptance additionally checks the input rows in Chromium.
-const source = readFileSync(new URL("../distributedValue.js", import.meta.url), "utf8")
-    .replace(/^import .*;\r?\n/gm, "");
-
-function setup({ connected = false, targetType = "STRING" } = {}) {
-    let extension;
-    const listeners = new Map();
-    const tasks = [];
-    const workers = [
-        { id: "gpu-a", name: "GPU A", enabled: true },
-        { id: "gpu-b", name: "GPU B", enabled: true },
-        { id: "off", name: "Disabled", enabled: false },
-    ];
-    const target = {
-        inputs: [{ name: "value" }],
-        widgets: [{ name: "value", type: targetType === "INT" ? "number" : "string", options: { step: 1, precision: 0 } }],
-    };
-    const graph = { links: { 1: { target_id: 2, target_slot: 0 } }, getNodeById: () => target };
-    const raw = ["default_value", "worker_values"].map((name, index) => ({
-        name,
-        type: "text",
-        value: index === 0 ? "saved default" : '{"1":"saved worker"}',
-        options: { multiline: false },
-        computeSize: vi.fn(() => [200, 20]),
-        serializeValue: vi.fn(function () { return this.value; }),
-    }));
-    const node = {
-        comfyClass: "DistributedValue", graph, widgets: [...raw], size: [200, 100],
-        inputs: raw.map(w => ({ name: w.name, widget: { name: w.name }, link: null })),
-        outputs: [{ links: connected ? [1] : [] }],
-        computeSize: () => [200, 100], setSize: vi.fn(), setDirtyCanvas: vi.fn(),
-        addWidget(type, name, value, callback, options) {
-            const widget = { type, name, value, callback, options };
-            this.widgets.push(widget);
-            return widget;
-        },
-        configure(data) {
-            data.widgets_values.forEach((value, i) => { this.widgets[i].value = value; });
-            return "configured";
-        },
-    };
-    vm.runInNewContext(source, {
-        app: { graph, registerExtension: ext => { extension = ext; } },
-        ENDPOINTS: { CONFIG: "/distributed/config" },
-        fetch: async () => ({ ok: true, json: async () => ({ workers }) }),
-        window: { addEventListener: (name, handler) => listeners.set(name, handler) },
-        setTimeout: callback => { tasks.push(callback); }, console,
-    }, { filename: "distributedValue.js" });
-    return { node, raw, extension, listeners, tasks, flush: () => { while (tasks.length) tasks.shift()(); } };
-}
-
-describe("Distributed Value internal widget visibility", () => {
-    it("hides raw fields natively without converting their types or replacing their options", async () => {
-        const h = setup();
-        const options = h.raw.map(w => w.options);
-        const sizes = h.raw.map(w => w.computeSize);
-        const serializers = h.raw.map(w => w.serializeValue);
-        const slots = h.node.inputs;
-        await h.extension.nodeCreated(h.node);
-
-        for (const [i, widget] of h.raw.entries()) {
-            expect(widget.type).toBe("text");
-            expect(widget.options).toBe(options[i]);
-            expect(widget.options.hidden).toBe(true);
-            expect(widget.computeSize).toBe(sizes[i]);
-            expect(widget.serializeValue).toBe(serializers[i]);
-            expect(h.node.widgets[i]).toBe(widget);
-        }
-        expect(h.node.inputs).toBe(slots);
-        expect(h.node.widgets.slice(2).map(w => w.label)).toEqual(["default_value", "GPU A", "GPU B"]);
-        expect(h.node.widgets.slice(2).every(w => !w.options.hidden)).toBe(true);
-    });
-
-    it("keeps typed edits in the original serialized fields across rebuild/configure", async () => {
-        const h = setup({ connected: true, targetType: "INT" });
-        await h.extension.nodeCreated(h.node);
-        h.node.widgets.find(w => w.name === "_dv_default").callback(42);
-        h.node.widgets.find(w => w.name === "_dv_worker_1").callback(7);
-        h.node.widgets.find(w => w.name === "_dv_worker_2").callback(9);
-        const values = h.raw.map(w => w.serializeValue());
-        expect(values[0]).toBe(42);
-        expect(JSON.parse(values[1])).toMatchObject({
-            _type: "INT", 1: "7", 2: "9", _by_worker_id: { "gpu-a": "7", "gpu-b": "9" },
-        });
-        expect(h.node.configure({ widgets_values: values })).toBe("configured");
-        h.flush();
-        expect(h.raw.map(w => w.serializeValue())).toEqual(values);
-        expect(h.raw.every(w => w.options.hidden && w.type === "text")).toBe(true);
-        expect(h.node.widgets.slice(2).map(w => w.value)).toEqual([42, 7, 9]);
-    });
-
-    it("reapplies visibility on worker refresh without duplicating raw or dynamic fields", async () => {
-        const h = setup();
-        await h.extension.nodeCreated(h.node);
-        h.raw[0].options.hidden = false;
-        await h.listeners.get("distributed:workers-changed")({ detail: { workers: [
-            { id: "gpu-b", name: "Renamed B", enabled: true },
-        ] } });
-        expect(h.raw.every(w => w.options.hidden)).toBe(true);
-        expect(h.node.widgets.map(w => w.name)).toEqual(["default_value", "worker_values", "_dv_default", "_dv_worker_1"]);
-        expect(h.node.widgets[3].label).toBe("Renamed B");
-    });
-
-    it("hides linked backing widgets and handles absent options without altering serializers", async () => {
-        const h = setup();
-        delete h.raw[1].options;
-        const linked = { type: "text", options: {}, serializeValue: vi.fn(() => "linked") };
-        h.raw[0].linkedWidgets = [linked];
-        await h.extension.nodeCreated(h.node);
-        expect(h.raw[1].options.hidden).toBe(true);
-        expect(linked.options.hidden).toBe(true);
-        expect(linked.type).toBe("text");
-        expect(linked.serializeValue()).toBe("linked");
-    });
-});

web/tests/executionUtils.test.js

@@ -0,0 +1,15 @@
+import { describe, expect, it } from "vitest";
+
+import { buildWorkerWebSocketUrl } from "../urlUtils.js";
+
+
+describe("execution decision helpers", () => {
+    it("buildWorkerWebSocketUrl converts http/https to ws/wss", () => {
+        expect(buildWorkerWebSocketUrl("http://worker.local:8188")).toBe(
+            "ws://worker.local:8188/distributed/worker_ws"
+        );
+        expect(buildWorkerWebSocketUrl("https://worker.example.com")).toBe(
+            "wss://worker.example.com/distributed/worker_ws"
+        );
+    });
+});

web/tests/imageListExecution.test.js

@@ -1,83 +0,0 @@
-import { afterEach, describe, expect, it, vi } from 'vitest';
-vi.mock('../../scripts/api.js', () => ({ api: { clientId: 'test-client' } }), { virtual: true });
-vi.mock('../workerLifecycle.js', () => ({
-    checkAllWorkerStatuses: vi.fn(), getWorkerUrl: () => 'http://worker.invalid',
-}));
-vi.mock('../workerUtils.js', () => ({
-    applyProbeResultToWorkerDot: vi.fn(),
-    findNodesByClass: (prompt, cls) => Object.values(prompt).filter(node => node.class_type === cls),
-}));
-import { api } from '../../scripts/api.js';
-import { setupInterceptor } from '../executionUtils.js';
-
-const prompt = { output: { '1': { class_type: 'DistributedImageListStart', inputs: {} } } };
-function extension(workers, online = true) {
-    return {
-        isEnabled: true, enabledWorkers: workers, log: vi.fn(), config: { settings: {} },
-        originalQueuePrompt: vi.fn().mockResolvedValue({ prompt_id: 'local' }),
-        api: { probeWorker: vi.fn().mockResolvedValue({ ok: online }),
-            getServerlessStatus: vi.fn(),
-            queueDistributed: vi.fn().mockResolvedValue({ prompt_id: 'distributed' }) },
-        ui: { showToast: vi.fn() },
-    };
-}
-afterEach(() => vi.useRealTimers());
-describe('image-list queue interception', () => {
-    it('routes image-list graphs through backend orchestration', async () => {
-        vi.useFakeTimers();
-        const ext = extension([{ id: 'one', type: 'local' }]);
-        setupInterceptor(ext);
-        expect(await api.queuePrompt(0, prompt)).toEqual({ prompt_id: 'distributed' });
-        expect(ext.api.queueDistributed).toHaveBeenCalledWith(expect.objectContaining({ enabled_worker_ids: ['one'] }));
-        expect(ext.originalQueuePrompt).not.toHaveBeenCalled();
-    });
-    it('does not fall back when orchestrator-only workers are offline', async () => {
-        vi.useFakeTimers();
-        const ext = extension([{ id: 'one', type: 'local' }], false);
-        ext.config.settings.master_delegate_only = true;
-        setupInterceptor(ext);
-        expect(await api.queuePrompt(0, prompt)).toBeNull();
-        expect(ext.originalQueuePrompt).not.toHaveBeenCalled();
-        expect(ext.api.queueDistributed).not.toHaveBeenCalled();
-    });
-    it('does not probe serverless providers for an image-list job', async () => {
-        vi.useFakeTimers();
-        const ext = extension([{ id: 'pool', type: 'serverless' }]);
-        setupInterceptor(ext);
-        expect(await api.queuePrompt(0, prompt)).toEqual({ prompt_id: 'distributed' });
-        expect(ext.api.queueDistributed).toHaveBeenCalledWith(expect.objectContaining({ enabled_worker_ids: [], delegate_master: false }));
-        expect(ext.api.getServerlessStatus).not.toHaveBeenCalled();
-        expect(ext.originalQueuePrompt).not.toHaveBeenCalled();
-    });
-    it.each([{ workers: [] }, { workers: [{ id: 'one', type: 'local' }] }])('queues master-only through orchestration when no workers are online: $workers', async ({ workers }) => {
-        vi.useFakeTimers();
-        const ext = extension(workers, false);
-        setupInterceptor(ext);
-        expect(await api.queuePrompt(0, prompt)).toEqual({ prompt_id: 'distributed' });
-        expect(ext.api.queueDistributed).toHaveBeenCalledWith(expect.objectContaining({ enabled_worker_ids: [], delegate_master: false }));
-        expect(ext.originalQueuePrompt).not.toHaveBeenCalled();
-    });
-    it('passes orchestrator-only mode to backend when workers are online', async () => {
-        vi.useFakeTimers();
-        const ext = extension([{ id: 'one', type: 'local' }]);
-        ext.config.settings.master_delegate_only = true;
-        setupInterceptor(ext);
-        await api.queuePrompt(0, prompt);
-        expect(ext.api.queueDistributed).toHaveBeenCalledWith(expect.objectContaining({ delegate_master: true }));
-    });
-    it('does not bypass backend errors with local execution', async () => {
-        vi.useFakeTimers();
-        const ext = extension([]);
-        ext.api.queueDistributed.mockRejectedValue(new Error('invalid region'));
-        setupInterceptor(ext);
-        expect(await api.queuePrompt(0, prompt)).toBeNull();
-        expect(ext.api.queueDistributed).toHaveBeenCalled();
-        expect(ext.originalQueuePrompt).not.toHaveBeenCalled();
-    });
-    it('preserves ordinary prompt handling', async () => {
-        const ext = extension([]);
-        setupInterceptor(ext);
-        expect(await api.queuePrompt(0, { output: {} })).toEqual({ prompt_id: 'local' });
-        expect(ext.api.queueDistributed).not.toHaveBeenCalled();
-    });
-});

web/tests/imageListStart.test.js

@@ -1,174 +0,0 @@
-import { readFileSync } from "node:fs";
-import vm from "node:vm";
-import { describe, expect, it } from "vitest";
-
-const source = readFileSync(new URL("../imageListStart.js", import.meta.url), "utf8")
-    .replace(/^import .*;\r?\n/gm, "");
-
-// Minimal LiteGraph-like graph: enough link bookkeeping to exercise output
-// add/remove, retyping, and link moves the extension performs.
-function setup() {
-    let extension;
-    const tasks = [];
-    const links = new Map();
-    let nextLink = 1;
-    const nodes = new Map();
-    const graph = {
-        links,
-        getLink: id => links.get(id),
-        getNodeById: id => nodes.get(id),
-    };
-    const makeSource = (id, type) => {
-        const node = { id, outputs: [{ type, links: [] }] };
-        nodes.set(id, node);
-        return node;
-    };
-    const makeTarget = (id, type) => {
-        const node = {
-            id, inputs: [{ type, link: null }],
-            disconnectInput(slot) {
-                const link = links.get(this.inputs[slot].link);
-                if (!link) return;
-                const origin = nodes.get(link.origin_id);
-                origin.outputs[link.origin_slot].links =
-                    origin.outputs[link.origin_slot].links.filter(l => l !== link.id);
-                links.delete(link.id);
-                this.inputs[slot].link = null;
-            },
-        };
-        nodes.set(id, node);
-        return node;
-    };
-    const fixed = ["image", "mask", "item_index", "context"].map(name => ({ name, type: "IMAGE", links: [], shape: 6 }));
-    const start = {
-        id: 1, comfyClass: "DistributedImageListStart", graph,
-        inputs: [{ name: "images", link: null }, { name: "mask", link: null }, { name: "values.value_1", link: null }],
-        outputs: [...fixed, ...Array.from({ length: 10 }, (_, i) => ({ name: `value_${i + 1}`, type: "*", links: [] }))],
-        addOutput(name, type, extra) { this.outputs.push({ name, type, links: [], ...(extra ?? {}) }); },
-        removeOutput(slot) {
-            if (this.outputs[slot].links.length) throw new Error("removed a wired output");
-            this.outputs.splice(slot, 1);
-        },
-        connect(slot, target, targetSlot) {
-            // LiteGraph refuses a link whose output type does not fit the input.
-            const outType = this.outputs[slot].type;
-            const inType = target.inputs[targetSlot].type;
-            if (outType !== "*" && inType !== "*" && !String(inType).split(",").includes(outType)) return null;
-            target.disconnectInput(targetSlot);
-            const id = nextLink++;
-            links.set(id, { id, origin_id: this.id, origin_slot: slot, target_id: target.id,
-                            target_slot: targetSlot, type: this.outputs[slot].type });
-            this.outputs[slot].links.push(id);
-            target.inputs[targetSlot].link = id;
-            return links.get(id);
-        },
-        disconnectOutput(slot) {
-            for (const id of [...this.outputs[slot].links]) {
-                const link = links.get(id);
-                nodes.get(link.target_id).disconnectInput(link.target_slot);
-            }
-        },
-        setDirtyCanvas() {},
-    };
-    nodes.set(1, start);
-    vm.runInNewContext(source, {
-        app: { registerExtension: ext => { extension = ext; } },
-        setTimeout: callback => { tasks.push(callback); },
-        console,
-    }, { filename: "imageListStart.js" });
-    extension.nodeCreated(start);
-    const flush = () => { while (tasks.length) tasks.shift()(); };
-    // Wire a source output into Start's value input, as Autogrow would.
-    const connectValue = (ordinal, sourceNode) => {
-        const name = `values.value_${ordinal}`;
-        let input = start.inputs.find(i => i.name === name);
-        if (!input) { input = { name, link: null }; start.inputs.push(input); }
-        const id = nextLink++;
-        links.set(id, { id, origin_id: sourceNode.id, origin_slot: 0, target_id: 1,
-                        target_slot: start.inputs.indexOf(input), type: sourceNode.outputs[0].type });
-        sourceNode.outputs[0].links.push(id);
-        input.link = id;
-        if (!start.inputs.some(i => i.name === `values.value_${ordinal + 1}`)) {
-            start.inputs.push({ name: `values.value_${ordinal + 1}`, link: null });
-        }
-        start.onConnectionsChange(1, start.inputs.indexOf(input), true);
-        flush();
-    };
-    // Mimic Autogrow's compaction after disconnecting a middle input.
-    const disconnectValueAndCompact = ordinal => {
-        const valueInputs = start.inputs.filter(i => i.name.startsWith("values."));
-        const linksBefore = valueInputs.map(i => i.link);
-        links.delete(linksBefore[ordinal - 1]);
-        const remaining = linksBefore.filter((l, i) => i !== ordinal - 1 && l != null);
-        start.inputs = start.inputs.filter(i => !i.name.startsWith("values."));
-        remaining.forEach((l, i) => start.inputs.push({ name: `values.value_${i + 1}`, link: l }));
-        start.inputs.push({ name: `values.value_${remaining.length + 1}`, link: null });
-        start.onConnectionsChange(1, 0, false);
-        flush();
-    };
-    flush();
-    return { start, makeSource, makeTarget, connectValue, disconnectValueAndCompact, links, flush };
-}
-
-const valueOutputs = start => start.outputs.slice(4).map(o => `${o.name}:${o.type}`);
-
-describe("Distributed Image List Start value outputs", () => {
-    it("shows no value outputs until a value input is connected", () => {
-        const h = setup();
-        expect(h.start.outputs.map(o => o.name)).toEqual(["image", "mask", "item_index", "context"]);
-    });
-
-    it("adds one typed output per connected value input", () => {
-        const h = setup();
-        h.connectValue(1, h.makeSource(10, "IMAGE"));
-        h.connectValue(2, h.makeSource(11, "STRING"));
-        expect(valueOutputs(h.start)).toEqual(["value_1:IMAGE", "value_2:STRING"]);
-        expect(h.start.outputs[4].shape).toBe(6);
-    });
-
-    it("keeps downstream wiring on the same value when a middle input is removed", () => {
-        const h = setup();
-        h.connectValue(1, h.makeSource(10, "IMAGE"));
-        h.connectValue(2, h.makeSource(11, "STRING"));
-        h.connectValue(3, h.makeSource(12, "INT"));
-        const imageUser = h.makeTarget(20, "IMAGE");
-        const stringUser = h.makeTarget(21, "STRING");
-        const intUser = h.makeTarget(22, "INT");
-        h.start.connect(4, imageUser, 0);
-        h.start.connect(5, stringUser, 0);
-        h.start.connect(6, intUser, 0);
-
-        h.disconnectValueAndCompact(2);
-
-        expect(valueOutputs(h.start)).toEqual(["value_1:IMAGE", "value_2:INT"]);
-        expect(h.links.get(imageUser.inputs[0].link).origin_slot).toBe(4);
-        expect(h.links.get(intUser.inputs[0].link).origin_slot).toBe(5);
-        expect(stringUser.inputs[0].link).toBe(null);
-    });
-
-    it("disconnects a consumer whose type no longer matches the input", () => {
-        const h = setup();
-        h.connectValue(1, h.makeSource(10, "IMAGE"));
-        const imageUser = h.makeTarget(20, "IMAGE");
-        h.start.connect(4, imageUser, 0);
-        // Replace value_1's source with a STRING source.
-        const input = h.start.inputs.find(i => i.name === "values.value_1");
-        h.links.delete(input.link);
-        input.link = null;
-        h.connectValue(1, h.makeSource(11, "STRING"));
-        expect(valueOutputs(h.start)).toEqual(["value_1:STRING"]);
-        expect(imageUser.inputs[0].link).toBe(null);
-    });
-
-    it("never removes an output that still has links", () => {
-        const h = setup();
-        h.connectValue(1, h.makeSource(10, "IMAGE"));
-        h.start.connect(4, h.makeTarget(20, "*"), 0);
-        const input = h.start.inputs.find(i => i.name === "values.value_1");
-        h.links.delete(input.link);
-        input.link = null;
-        h.start.onConnectionsChange(1, 2, false);
-        h.flush();
-        expect(valueOutputs(h.start)).toEqual(["value_1:*"]);
-    });
-});

web/tests/masterDetection.test.js

@@ -1,65 +0,0 @@
-import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
-import { detectMasterIP } from "../masterDetection.js";
-
-describe("automatic local worker ports", () => {
-    let originalWindow;
-    beforeEach(() => {
-        originalWindow = globalThis.window;
-        globalThis.window = { location: { hostname: "localhost", port: "443" } };
-    });
-    afterEach(() => { globalThis.window = originalWindow; });
-
-    function makeExtension(overrides = {}) {
-        return {
-            config: { workers: [], settings: {}, master: {} },
-            log: vi.fn(),
-            api: {
-                getNetworkInfo: vi.fn().mockResolvedValue({
-                    cuda_device: 0, cuda_device_count: 3,
-                    master_port: 8189, local_worker_ports: [8191, 8193],
-                    ...overrides,
-                }),
-                updateMaster: vi.fn().mockResolvedValue({}),
-                updateWorker: vi.fn().mockResolvedValue({}),
-                updateSetting: vi.fn().mockResolvedValue({}),
-            },
-            ui: { updateMasterDisplay: vi.fn() },
-            app: {},
-            loadConfig: vi.fn().mockResolvedValue(),
-        };
-    }
-
-    it("uses server-allocated ports, not hardcoded ports or the browser proxy port", async () => {
-        const extension = makeExtension();
-        await detectMasterIP(extension);
-        expect(extension.config.workers.map(worker => worker.port)).toEqual([8191, 8193]);
-        expect(extension.config.workers.map(worker => worker.cuda_device)).toEqual([1, 2]);
-        expect(extension.api.updateWorker).toHaveBeenCalledTimes(2);
-    });
-
-    it("does not overwrite an existing manually configured worker", async () => {
-        const extension = makeExtension();
-        const worker = { id: "manual", host: "localhost", port: 8189 };
-        extension.config.workers = [worker];
-        await detectMasterIP(extension);
-        expect(extension.config.workers).toEqual([worker]);
-        expect(extension.api.updateWorker).not.toHaveBeenCalled();
-    });
-
-    it("fails without partially saving workers if not enough ports are available", async () => {
-        const extension = makeExtension({ local_worker_ports: [8191] });
-        await detectMasterIP(extension);
-        expect(extension.config.workers).toEqual([]);
-        expect(extension.api.updateWorker).not.toHaveBeenCalled();
-        expect(extension.api.updateSetting).not.toHaveBeenCalled();
-        expect(extension.log).toHaveBeenCalledWith(expect.stringContaining("port"), "error");
-    });
-
-    it("preserves Runpod worker host and launch settings", async () => {
-        globalThis.window.location.hostname = "pod-8189.proxy.runpod.net";
-        const extension = makeExtension();
-        await detectMasterIP(extension);
-        expect(extension.config.workers.map(worker => worker.port)).toEqual([8191, 8193]);
-        expect(extension.config.workers.every(worker => worker.host === null && worker.extra_args === "--listen")).toBe(true);
-    });
-});

web/tests/serverlessWorkers.test.js

@@ -1,622 +0,0 @@
-import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
-
-vi.mock("../../scripts/api.js", () => ({ api: { clientId: "test-client" } }), { virtual: true });
-
-import { createApiClient } from "../apiClient.js";
-import { createWorkerSettingsForm } from "../ui/settingsForm.js";
-import { DistributedUI } from "../ui.js";
-import { checkWorkerStatus, formatServerlessStatus, getWorkerUrl, launchWorker, stopWorker, viewWorkerLog } from "../workerLifecycle.js";
-import { handleClearMemory, handleInterruptWorkers } from "../workerUtils.js";
-import { isRemoteWorker, normalizeRunpodEndpointId, saveWorkerSettings, testServerlessWorker } from "../workerSettings.js";
-import { performPreflightCheck } from "../executionUtils.js";
-
-function createDocumentMock() {
-    const elements = new Map();
-    const makeElement = (tagName) => {
-        const element = {
-            tagName,
-            children: [],
-            dataset: {},
-            style: {},
-            setAttribute(name, value) { this[name] = value; },
-            classList: {
-                values: new Set(),
-                add(...names) { names.forEach((name) => this.values.add(name)); },
-                remove(...names) { names.forEach((name) => this.values.delete(name)); },
-                toggle(name, force) {
-                    const enabled = force ?? !this.values.has(name);
-                    if (enabled) this.values.add(name);
-                    else this.values.delete(name);
-                    return enabled;
-                },
-            },
-            appendChild(child) { this.children.push(child); child.parentNode = this; return child; },
-            dispatchEvent() { this.onchange?.({ target: this }); },
-        };
-        Object.defineProperty(element, "id", {
-            get() { return this._id || ""; },
-            set(value) { this._id = value; elements.set(value, this); },
-        });
-        return element;
-    };
-    return {
-        elements,
-        document: {
-            createElement: makeElement,
-            getElementById: (id) => elements.get(id) || null,
-        },
-    };
-}
-
-function createFormUi() {
-    return {
-        createFormGroup(label, value, id, type = "text", placeholder = "") {
-            const group = document.createElement("div");
-            const input = document.createElement("input");
-            input.id = id;
-            input.value = String(value ?? "");
-            input.type = type;
-            input.placeholder = placeholder;
-            group.appendChild(input);
-            return { group, input, label };
-        },
-        createButton(text, onclick) {
-            const button = document.createElement("button");
-            button.textContent = text;
-            button.onclick = onclick;
-            return button;
-        },
-        createButtonGroup(buttons) {
-            const group = document.createElement("div");
-            buttons.forEach((button) => group.appendChild(button));
-            return group;
-        },
-    };
-}
-
-function makeSaveExtension(worker, updateWorker = vi.fn().mockResolvedValue({})) {
-    return {
-        config: { workers: [worker] },
-        api: { updateWorker },
-        state: { updateWorker: vi.fn(), setWorkerExpanded: vi.fn() },
-        app: { extensionManager: { toast: { add: vi.fn() } } },
-        panelElement: null,
-    };
-}
-
-describe("Runpod serverless worker panel", () => {
-    let originalDocument;
-    let originalWindow;
-    let elements;
-
-    beforeEach(() => {
-        originalDocument = globalThis.document;
-        originalWindow = globalThis.window;
-        const mock = createDocumentMock();
-        elements = mock.elements;
-        globalThis.document = mock.document;
-        globalThis.window = {
-            location: { hostname: "localhost", protocol: "http:", port: "8188", origin: "http://localhost:8188" },
-        };
-    });
-
-    afterEach(() => {
-        globalThis.document = originalDocument;
-        globalThis.window = originalWindow;
-        vi.restoreAllMocks();
-    });
-
-    it("classifies Runpod serverless as remote for worker allocation", () => {
-        expect(isRemoteWorker({}, { type: "serverless" })).toBe(true);
-    });
-
-    it("initializes a saved serverless form before it is attached to the document", () => {
-        // Real document lookups cannot see a form while its card is being built.
-        vi.spyOn(document, "getElementById").mockReturnValue(null);
-        const worker = {
-            id: "saved", name: "Runpod", type: "serverless", endpoint_id: "ep_123",
-            jobs_per_run: 2, concurrency: 1, timeout_seconds: 900, api_key_configured: true,
-        };
-        const form = createWorkerSettingsForm(createFormUi(), {}, worker);
-
-        expect(form.parentNode).toBeUndefined();
-        expect(elements.get("worker-type-saved").value).toBe("serverless");
-        for (const field of ["host", "port", "cuda", "args"]) {
-            expect(elements.get(`${field}-group-saved`).style.display).toBe("none");
-        }
-        expect(elements.get("runpod-text-saved").style.display).toBe("none");
-        for (const field of ["endpoint-id", "api-key", "clear-api-key", "jobs-per-run", "timeout-seconds", "serverless-test"]) {
-            expect(elements.get(`${field}-group-saved`).style.display).toBe("flex");
-        }
-        // The separate concurrency setting was removed; a legacy value is ignored.
-        expect(elements.has("concurrency-saved")).toBe(false);
-        expect(elements.has("concurrency-group-saved")).toBe(false);
-        expect(elements.get("endpoint-id-saved").value).toBe("ep_123");
-        expect(elements.get("jobs-per-run-saved").value).toBe("2");
-        expect(elements.get("timeout-seconds-saved").value).toBe("900");
-        expect(elements.get("api-key-saved").value).toBe("");
-        expect(elements.get("api-key-saved").placeholder).toMatch(/leave blank/i);
-    });
-
-    it("preserves a saved cloud port on first render and applies defaults only on type changes", () => {
-        vi.spyOn(document, "getElementById").mockReturnValue(null);
-        createWorkerSettingsForm(createFormUi(), {}, {
-            id: "cloud", name: "Cloud", type: "cloud", host: "worker.example", port: 8443,
-        });
-        const type = elements.get("worker-type-cloud");
-        expect(type.value).toBe("cloud");
-        expect(elements.get("port-cloud").value).toBe("8443");
-        expect(elements.get("host-cloud").value).toBe("worker.example");
-        expect(elements.get("runpod-text-cloud").style.display).toBe("block");
-
-        type.value = "remote";
-        type.dispatchEvent();
-        type.value = "cloud";
-        type.dispatchEvent();
-        expect(elements.get("port-cloud").value).toBe("443");
-    });
-
-    it("shows only serverless settings and keeps credentials write-only", () => {
-        const worker = {
-            id: "w1", name: "Runpod", type: "remote", host: "worker.example", port: 8188,
-            cuda_device: 0, extra_args: "--test", api_key_configured: true,
-        };
-        const form = createWorkerSettingsForm(createFormUi(), {}, worker);
-        const type = elements.get("worker-type-w1");
-        type.value = "serverless";
-        type.dispatchEvent();
-
-        expect(type.children.map((option) => [option.value, option.textContent])).toContainEqual([
-            "serverless", "Serverless (Runpod)",
-        ]);
-        for (const id of ["host-group-w1", "port-group-w1", "cuda-group-w1", "args-group-w1"]) {
-            expect(elements.get(id).style.display).toBe("none");
-        }
-        for (const id of ["endpoint-id-w1", "api-key-w1", "jobs-per-run-w1", "timeout-seconds-w1"]) {
-            expect(elements.has(id)).toBe(true);
-        }
-        expect(elements.has("concurrency-w1")).toBe(false);
-        for (const id of ["endpoint-id-group-w1", "api-key-group-w1", "clear-api-key-group-w1", "jobs-per-run-group-w1", "timeout-seconds-group-w1", "serverless-test-group-w1"]) {
-            expect(elements.get(id).style.display).toBe("flex");
-        }
-        expect(elements.has("serverless-help-w1")).toBe(false);
-        expect(elements.get("api-key-w1").type).toBe("password");
-        expect(elements.get("api-key-w1").value).toBe("");
-        expect(elements.get("api-key-w1").placeholder).toMatch(/leave blank/i);
-        expect(elements.get("jobs-per-run-w1").value).toBe("1");
-        expect(elements.get("jobs-per-run-w1").min).toBe("1");
-        expect(elements.get("jobs-per-run-w1").max).toBe("16");
-        expect(elements.get("jobs-per-run-w1").title).toBe(
-            "Separate Runpod jobs submitted each run, all at once. Runpod's max workers decides how many GPUs actually run them.",
-        );
-        expect(elements.get("timeout-seconds-w1").value).toBe("1800");
-        expect(elements.get("timeout-seconds-w1").min).toBe("30");
-        expect(elements.get("timeout-seconds-w1").max).toBe("86400");
-        expect(form.children.length).toBeGreaterThan(0);
-    });
-
-    it("disables Clear saved API key while the worker is enabled and explains why", () => {
-        vi.spyOn(document, "getElementById").mockReturnValue(null);
-        const base = { name: "Runpod", type: "serverless", endpoint_id: "ep_123" };
-        createWorkerSettingsForm(createFormUi(), {}, { ...base, id: "on", enabled: true, api_key_configured: true });
-        createWorkerSettingsForm(createFormUi(), {}, { ...base, id: "off", enabled: false, api_key_configured: true });
-        createWorkerSettingsForm(createFormUi(), {}, { ...base, id: "nokey", enabled: false, api_key_configured: false });
-
-        expect(elements.get("clear-api-key-on").disabled).toBe(true);
-        expect(elements.get("clear-api-key-group-on").title).toMatch(/turn this worker off/i);
-        expect(elements.get("clear-api-key-off").disabled).toBe(false);
-        expect(elements.get("clear-api-key-nokey").disabled).toBe(true);
-        expect(elements.get("clear-api-key-group-nokey").title).toMatch(/no api key is saved/i);
-        for (const id of ["on", "off", "nokey"]) {
-            expect(elements.get(`clear-api-key-group-${id}`).title).not.toMatch(/pool/i);
-        }
-    });
-
-    it("saves fixed Runpod settings without enabling or retaining the API key", async () => {
-        const worker = { id: "w2", name: "Old", type: "remote", enabled: false, api_key_configured: false };
-        const extension = makeSaveExtension(worker);
-        elements.set("name-w2", { value: "Runpod worker" });
-        elements.set("worker-type-w2", { value: "serverless" });
-        elements.set("endpoint-id-w2", { value: "ep_123" });
-        elements.set("api-key-w2", { value: "secret-token" });
-        elements.set("clear-api-key-w2", { checked: false });
-        elements.set("jobs-per-run-w2", { value: "4" });
-        elements.set("timeout-seconds-w2", { value: "2400" });
-
-        await saveWorkerSettings(extension, "w2");
-
-        expect(extension.api.updateWorker).toHaveBeenCalledWith("w2", {
-            name: "Runpod worker", type: "serverless", provider: "runpod", endpoint_id: "ep_123",
-            api_key: "secret-token", jobs_per_run: 4, timeout_seconds: 2400, enabled: false,
-        });
-        expect(worker).toMatchObject({ type: "serverless", provider: "runpod", enabled: false, api_key_configured: true });
-        expect(worker).not.toHaveProperty("api_key");
-        expect(JSON.stringify(extension.config)).not.toContain("secret-token");
-        expect(elements.get("api-key-w2").value).toBe("");
-        expect(extension.state.updateWorker).toHaveBeenCalledWith("w2", { enabled: false });
-        expect(JSON.stringify(extension.state.updateWorker.mock.calls)).not.toContain("secret-token");
-    });
-
-    it("labels the single parallelism setting Images per run and drops a legacy concurrency value", async () => {
-        const worker = { id: "legacy", name: "Runpod", type: "serverless", provider: "runpod", enabled: false,
-            api_key_configured: true, jobs_per_run: 1, concurrency: 1 };
-        const extension = makeSaveExtension(worker);
-        elements.set("name-legacy", { value: "Runpod" });
-        elements.set("worker-type-legacy", { value: "serverless" });
-        elements.set("endpoint-id-legacy", { value: "ep_123" });
-        elements.set("api-key-legacy", { value: "" });
-        elements.set("clear-api-key-legacy", { checked: false });
-        elements.set("jobs-per-run-legacy", { value: "2" });
-        elements.set("timeout-seconds-legacy", { value: "1800" });
-
-        await saveWorkerSettings(extension, "legacy");
-
-        const payload = extension.api.updateWorker.mock.calls[0][1];
-        expect(payload.jobs_per_run).toBe(2);
-        expect(payload).not.toHaveProperty("concurrency");
-        expect(worker.jobs_per_run).toBe(2);
-        expect(worker).not.toHaveProperty("concurrency");
-
-        elements.set("jobs-per-run-legacy", { value: "17" });
-        await saveWorkerSettings(extension, "legacy");
-        expect(extension.api.updateWorker).toHaveBeenCalledTimes(1);
-        expect(extension.app.extensionManager.toast.add).toHaveBeenLastCalledWith(
-            expect.objectContaining({ detail: "Images per run must be a whole number from 1 to 16" }),
-        );
-    });
-
-    it("trims a pasted Runpod URL to the endpoint ID and explains other invalid IDs", async () => {
-        expect(normalizeRunpodEndpointId("  abc123xyz ")).toBe("abc123xyz");
-        expect(normalizeRunpodEndpointId("https://api.runpod.ai/v2/abc123xyz/run")).toBe("abc123xyz");
-        expect(normalizeRunpodEndpointId("api.runpod.ai/v2/abc_12-3")).toBe("abc_12-3");
-        expect(normalizeRunpodEndpointId("https://www.runpod.io/console/serverless/user/endpoint/abc")).toBe(
-            "https://www.runpod.io/console/serverless/user/endpoint/abc",
-        );
-
-        const worker = { id: "url", name: "Runpod", type: "serverless", provider: "runpod", enabled: false, api_key_configured: true };
-        const extension = makeSaveExtension(worker);
-        elements.set("name-url", { value: "Runpod" });
-        elements.set("worker-type-url", { value: "serverless" });
-        elements.set("endpoint-id-url", { value: "https://api.runpod.ai/v2/abc123xyz/runsync" });
-        elements.set("api-key-url", { value: "" });
-        elements.set("clear-api-key-url", { checked: false });
-        elements.set("jobs-per-run-url", { value: "1" });
-        elements.set("timeout-seconds-url", { value: "1800" });
-        await saveWorkerSettings(extension, "url");
-        expect(extension.api.updateWorker.mock.calls[0][1].endpoint_id).toBe("abc123xyz");
-
-        elements.set("endpoint-id-url", { value: "https://www.runpod.io/console/serverless/endpoint/abc" });
-        await saveWorkerSettings(extension, "url");
-        expect(extension.api.updateWorker).toHaveBeenCalledTimes(1);
-        expect(extension.app.extensionManager.toast.add).toHaveBeenLastCalledWith(
-            expect.objectContaining({ detail: expect.stringMatching(/just the Runpod endpoint ID/) }),
-        );
-    });
-
-    it("asks before a type change deletes the saved API key, and cancels cleanly", async () => {
-        const worker = { id: "conv", name: "Runpod", type: "serverless", provider: "runpod", enabled: false, api_key_configured: true };
-        const extension = makeSaveExtension(worker);
-        elements.set("name-conv", { value: "Runpod" });
-        elements.set("worker-type-conv", { value: "remote" });
-        elements.set("host-conv", { value: "192.168.1.20" });
-        elements.set("port-conv", { value: "8188" });
-        const confirmSpy = vi.fn().mockReturnValue(false);
-        globalThis.confirm = confirmSpy;
-        try {
-            await saveWorkerSettings(extension, "conv");
-            expect(confirmSpy).toHaveBeenCalledWith(expect.stringMatching(/saved Runpod API key will be deleted/));
-            expect(extension.api.updateWorker).not.toHaveBeenCalled();
-            expect(worker.type).toBe("serverless");
-
-            confirmSpy.mockReturnValue(true);
-            await saveWorkerSettings(extension, "conv");
-            expect(extension.api.updateWorker).toHaveBeenCalledTimes(1);
-            expect(extension.api.updateWorker.mock.calls[0][1].type).toBe("remote");
-        } finally {
-            delete globalThis.confirm;
-        }
-    });
-
-    it("does not ask for confirmation when a serverless worker without a key changes type", async () => {
-        const worker = { id: "nokey2", name: "Runpod", type: "serverless", provider: "runpod", enabled: false, api_key_configured: false };
-        const extension = makeSaveExtension(worker);
-        elements.set("name-nokey2", { value: "Runpod" });
-        elements.set("worker-type-nokey2", { value: "remote" });
-        elements.set("host-nokey2", { value: "192.168.1.20" });
-        elements.set("port-nokey2", { value: "8188" });
-        const confirmSpy = vi.fn().mockReturnValue(false);
-        globalThis.confirm = confirmSpy;
-        try {
-            await saveWorkerSettings(extension, "nokey2");
-            expect(confirmSpy).not.toHaveBeenCalled();
-            expect(extension.api.updateWorker).toHaveBeenCalledTimes(1);
-        } finally {
-            delete globalThis.confirm;
-        }
-    });
-
-    it("omits a blank API key so the saved key is preserved", async () => {
-        const worker = { id: "w-blank", name: "Runpod", type: "serverless", provider: "runpod", enabled: true, api_key_configured: true };
-        const extension = makeSaveExtension(worker);
-        elements.set("name-w-blank", { value: "Runpod" });
-        elements.set("worker-type-w-blank", { value: "serverless" });
-        elements.set("endpoint-id-w-blank", { value: "ep_123" });
-        elements.set("api-key-w-blank", { value: "" });
-        elements.set("clear-api-key-w-blank", { checked: false });
-        elements.set("jobs-per-run-w-blank", { value: "1" });
-        elements.set("timeout-seconds-w-blank", { value: "1800" });
-
-        await saveWorkerSettings(extension, "w-blank");
-
-        const payload = extension.api.updateWorker.mock.calls[0][1];
-        expect(payload).not.toHaveProperty("api_key");
-        expect(payload).not.toHaveProperty("clear_api_key");
-        expect(worker.api_key_configured).toBe(true);
-        expect(worker).not.toHaveProperty("api_key");
-    });
-
-    it("omits a blank key to preserve it and sends an explicit clear flag when requested", async () => {
-        const worker = { id: "w3", name: "Runpod", type: "serverless", provider: "runpod", enabled: true, api_key_configured: true };
-        const extension = makeSaveExtension(worker);
-        elements.set("name-w3", { value: "Runpod" });
-        elements.set("worker-type-w3", { value: "serverless" });
-        elements.set("endpoint-id-w3", { value: "ep_123" });
-        elements.set("api-key-w3", { value: "" });
-        elements.set("clear-api-key-w3", { checked: true });
-        elements.set("jobs-per-run-w3", { value: "1" });
-        elements.set("timeout-seconds-w3", { value: "1800" });
-
-        await saveWorkerSettings(extension, "w3");
-
-        expect(extension.api.updateWorker).toHaveBeenCalledWith("w3", {
-            name: "Runpod", type: "serverless", provider: "runpod", endpoint_id: "ep_123",
-            clear_api_key: true, jobs_per_run: 1, timeout_seconds: 1800, enabled: true,
-        });
-        expect(worker.api_key_configured).toBe(false);
-        expect(worker).not.toHaveProperty("api_key");
-    });
-
-    it.each([
-        [undefined, "Checking status…"],
-        [{ ok: true, queueRemaining: 0, jobs: [] }, "Ready"],
-        [{ ok: true, queueRemaining: 1, jobs: [{ status: "IN_QUEUE" }] }, "Queued"],
-        [{ ok: true, jobs: [{ status: "IN_PROGRESS" }] }, "Running"],
-        [{ ok: true, jobs: [{ status: "RUNNING" }, { status: "IN_QUEUE" }] }, "Running"],
-        [{ ok: true, jobs: [{ status: "RUNNING" }, { status: "IN_PROGRESS" }] }, "Running · 2 jobs"],
-        [{ ok: true, jobs: [{ status: "DOWNLOADING" }] }, "Downloading result"],
-        [{ ok: true, jobs: [{ status: "CANCEL_REQUESTED" }] }, "Cancelling"],
-        [{ ok: true, queueRemaining: 1 }, "Active · 1 job"],
-        [{ ok: true, queueRemaining: 3 }, "Active · 3 jobs"],
-        [{ ok: true, jobs: [{ status: "COMPLETED" }] }, "Ready"],
-        [{ ok: true, jobs: [{ status: "FAILED" }] }, "Last job failed"],
-        [{ ok: false, message: "Check the saved endpoint ID and API key" }, "Unavailable"],
-    ])("formats provider state as one plain-language summary: %j", (status, expected) => {
-        expect(formatServerlessStatus(status)).toBe(expected);
-    });
-
-    it("keeps status above a consistent endpoint info box and retains expanded settings", () => {
-        const ui = new DistributedUI();
-        const worker = {
-            id: "compact", name: "Used Car Inventory", type: "serverless", enabled: true,
-            endpoint_id: "ep_private_id", jobs_per_run: 1, concurrency: 1,
-        };
-        const extension = {
-            config: { workers: [worker] },
-            state: {
-                isWorkerExpanded: () => false,
-                getWorker: () => ({ serverlessStatus: { ok: true, queueRemaining: 1, jobs: [{ status: "IN_QUEUE" }] } }),
-            },
-        };
-        const card = ui.renderEntityCard("worker", worker, extension);
-        const collect = (element) => [element, ...(element.children || []).flatMap(collect)];
-        const info = collect(card).find((element) => element.classList.values.has("dist-worker-info"));
-        expect(info.innerHTML).not.toContain("Runpod Serverless");
-        expect(info.innerHTML).not.toMatch(/ep_private_id|jobs\/run|slots/);
-        expect(elements.get("serverless-status-compact").textContent).toBe("Queued");
-        expect(elements.get("endpoint-id-compact").value).toBe("ep_private_id");
-        expect(Number(elements.get("jobs-per-run-compact").value)).toBe(1);
-        expect(elements.has("concurrency-compact")).toBe(false);
-        const box = elements.get("controls-compact").children[0].children[0];
-        expect(box.textContent).toBe("Runpod Serverless endpoint");
-        expect(box.classList.values.has("dist-info-box")).toBe(true);
-        expect(box.style.cssText).toBe(ui.styles.infoBox);
-        expect(elements.get("status-compact").style.cssText).toContain("flex: 0 0 10px");
-    });
-
-    it("uses provider status instead of a direct ComfyUI URL probe", async () => {
-        const statusElement = { textContent: "" };
-        elements.set("serverless-status-w4", statusElement);
-        const status = {
-            ok: true, queueRemaining: 0, message: "Ready to scale",
-            jobs: [{ job_id: "job-1", status: "RUNNING", worker_id: "w4", execution_id: "exec-1" }],
-        };
-        const extension = {
-            config: { workers: [{ id: "w4", name: "Runpod", type: "serverless", enabled: true }] },
-            api: { getServerlessStatus: vi.fn().mockResolvedValue(status), probeWorker: vi.fn() },
-            statusCheckAbortController: null,
-            state: {
-                setWorkerStatus: vi.fn(),
-                updateWorker: vi.fn(),
-                getWorkerStatus: vi.fn().mockReturnValue({ online: true, processing: false, queueCount: 0 }),
-                getWorker: vi.fn().mockReturnValue({}),
-                isWorkerLaunching: vi.fn().mockReturnValue(false),
-            },
-            updateWorkerCard: vi.fn().mockReturnValue(true),
-            log: vi.fn(),
-        };
-
-        await checkWorkerStatus(extension, extension.config.workers[0]);
-
-        expect(extension.api.getServerlessStatus).toHaveBeenCalledWith("w4");
-        expect(extension.api.probeWorker).not.toHaveBeenCalled();
-        expect(extension.state.setWorkerStatus).toHaveBeenCalledWith("w4", {
-            online: true, processing: true, queueCount: 0,
-        });
-        expect(statusElement.textContent).toBe("Running");
-        expect(statusElement.title).toBe("Ready to scale");
-    });
-
-    it("considers a healthy zero-worker Runpod pool ready and never builds a worker URL", async () => {
-        const worker = { id: "w-ready", name: "Runpod", type: "serverless", enabled: true };
-        const extension = {
-            api: {
-                getServerlessStatus: vi.fn().mockResolvedValue({ ok: true, queueRemaining: 0, message: "Ready to scale" }),
-                probeWorker: vi.fn(),
-            },
-            log: vi.fn(),
-        };
-
-        await expect(performPreflightCheck(extension, [worker])).resolves.toEqual([worker]);
-        expect(extension.api.getServerlessStatus).toHaveBeenCalledWith("w-ready");
-        expect(extension.api.probeWorker).not.toHaveBeenCalled();
-        expect(getWorkerUrl(extension, worker)).toBeNull();
-    });
-
-    it("does not render or invoke direct ComfyUI launch, stop, or log controls", async () => {
-        const ui = new DistributedUI();
-        const worker = { id: "w-controls", name: "Runpod", type: "serverless", enabled: true };
-        const controls = ui.createControlsSection({ dynamic: true }, worker, {}, true);
-
-        const extension = {
-            config: { workers: [worker] },
-            api: {
-                launchWorker: vi.fn(), stopWorker: vi.fn(), getRemoteWorkerLog: vi.fn(),
-            },
-        };
-
-        const box = controls.children[0].children[0];
-        expect(box.textContent).toBe("Runpod Serverless endpoint");
-        expect(box.tagName).toBe("div");
-        expect(controls.children[0].children).toHaveLength(1);
-        await launchWorker(extension, worker.id);
-        await stopWorker(extension, worker.id);
-        await viewWorkerLog(extension, worker.id, true);
-        expect(extension.api.launchWorker).not.toHaveBeenCalled();
-        expect(extension.api.stopWorker).not.toHaveBeenCalled();
-        expect(extension.api.getRemoteWorkerLog).not.toHaveBeenCalled();
-    });
-
-    it("uses the serverless status, test, and cancel routes with the frozen request shapes", async () => {
-        const originalFetch = globalThis.fetch;
-        const statusPayload = {
-            ok: true,
-            queueRemaining: 0,
-            message: "Ready to scale",
-            jobs: [{ job_id: "job-1", status: "RUNNING", worker_id: "w5", execution_id: "exec-1" }],
-        };
-        globalThis.fetch = vi.fn()
-            .mockResolvedValueOnce({ ok: true, status: 200, json: vi.fn().mockResolvedValue({ ok: true, message: "Endpoint test passed" }) })
-            .mockResolvedValueOnce({ ok: true, status: 200, json: vi.fn().mockResolvedValue(statusPayload) })
-            .mockResolvedValueOnce({ ok: true, status: 200, json: vi.fn().mockResolvedValue({ ok: true }) });
-        try {
-            const client = createApiClient("http://localhost:8188");
-            await client.testServerlessWorker("w5");
-            await expect(client.getServerlessStatus("w5")).resolves.toEqual(statusPayload);
-            await client.cancelServerlessWorker("w5");
-
-            expect(globalThis.fetch).toHaveBeenNthCalledWith(1, "http://localhost:8188/distributed/serverless/test", expect.objectContaining({
-                method: "POST", body: JSON.stringify({ worker_id: "w5" }),
-            }));
-            expect(globalThis.fetch).toHaveBeenNthCalledWith(2, "http://localhost:8188/distributed/serverless/status/w5", expect.any(Object));
-            expect(globalThis.fetch).toHaveBeenNthCalledWith(3, "http://localhost:8188/distributed/serverless/cancel", expect.objectContaining({
-                method: "POST", body: JSON.stringify({ worker_id: "w5" }),
-            }));
-        } finally {
-            globalThis.fetch = originalFetch;
-        }
-    });
-
-    it("tests only the saved Runpod endpoint and reports the backend message", async () => {
-        const button = document.createElement("button");
-        button.id = "test-serverless-w6";
-        button.textContent = "Test Saved Endpoint";
-        const extension = {
-            api: {
-                testServerlessWorker: vi.fn().mockResolvedValue({ ok: true, message: "Endpoint reachable" }),
-                queueDistributed: vi.fn(),
-            },
-            app: { extensionManager: { toast: { add: vi.fn() } } },
-        };
-
-        await expect(testServerlessWorker(extension, "w6")).resolves.toEqual({ ok: true, message: "Endpoint reachable" });
-
-        expect(extension.api.testServerlessWorker).toHaveBeenCalledWith("w6");
-        expect(extension.api.queueDistributed).not.toHaveBeenCalled();
-        expect(extension.app.extensionManager.toast.add).toHaveBeenCalledWith(expect.objectContaining({
-            severity: "success", detail: "Endpoint reachable",
-        }));
-        expect(button.textContent).toBe("Test Saved Endpoint");
-        expect(button.disabled).toBe(false);
-    });
-
-    it("redacts API-key-bearing worker update errors from logs and user-visible exceptions", async () => {
-        const originalFetch = globalThis.fetch;
-        const logSpy = vi.spyOn(console, "log").mockImplementation(() => {});
-        const secret = "never-log-this-runpod-key";
-        globalThis.fetch = vi.fn().mockResolvedValue({
-            ok: false,
-            status: 400,
-            json: vi.fn().mockResolvedValue({ message: `Rejected key ${secret}` }),
-        });
-        try {
-            const client = createApiClient("http://localhost:8188");
-            await expect(client.updateWorker("w-secret", { api_key: secret })).rejects.toThrow("[redacted]");
-            expect(logSpy.mock.calls.flat().join(" ")).not.toContain(secret);
-        } finally {
-            globalThis.fetch = originalFetch;
-        }
-    });
-
-    it("never sends VRAM operations to serverless workers", async () => {
-        const originalFetch = globalThis.fetch;
-        globalThis.fetch = vi.fn();
-        vi.useFakeTimers();
-        const button = {
-            textContent: "Clear Worker VRAM", style: { cssText: "" }, disabled: false,
-            classList: { names: new Set(), add(...names) { names.forEach((name) => this.names.add(name)); }, remove(...names) { names.forEach((name) => this.names.delete(name)); } },
-        };
-        const extension = { enabledWorkers: [{ id: "s-vram", name: "Runpod", type: "serverless", enabled: true }] };
-        try {
-            await handleClearMemory(extension, button);
-            expect(globalThis.fetch).not.toHaveBeenCalled();
-        } finally {
-            vi.runOnlyPendingTimers();
-            vi.useRealTimers();
-            globalThis.fetch = originalFetch;
-        }
-    });
-
-    it("cancels each enabled serverless worker once without changing normal worker interrupts", async () => {
-        const originalFetch = globalThis.fetch;
-        globalThis.fetch = vi.fn().mockResolvedValue({ ok: true });
-        vi.useFakeTimers();
-        const button = {
-            textContent: "Interrupt Workers", style: { cssText: "" }, disabled: false,
-            classList: { names: new Set(), add(...names) { names.forEach((name) => this.names.add(name)); }, remove(...names) { names.forEach((name) => this.names.delete(name)); } },
-        };
-        const workers = [
-            { id: "s1", name: "Serverless 1", type: "serverless", enabled: true },
-            { id: "s2", name: "Serverless 2", type: "serverless", enabled: true },
-            { id: "s3", name: "Disabled", type: "serverless", enabled: false },
-            { id: "local", name: "Local", type: "local", enabled: true, port: 8189 },
-        ];
-        const extension = {
-            enabledWorkers: workers.filter((worker) => worker.enabled),
-            api: { cancelServerlessWorker: vi.fn().mockResolvedValue({ ok: true }) },
-            log: vi.fn(),
-        };
-        try {
-            await handleInterruptWorkers(extension, button);
-
-            expect(extension.api.cancelServerlessWorker).toHaveBeenCalledTimes(2);
-            expect(extension.api.cancelServerlessWorker).toHaveBeenNthCalledWith(1, "s1");
-            expect(extension.api.cancelServerlessWorker).toHaveBeenNthCalledWith(2, "s2");
-            expect(globalThis.fetch).toHaveBeenCalledTimes(1);
-            expect(globalThis.fetch.mock.calls[0][0]).toBe("http://localhost:8189/interrupt");
-        } finally {
-            vi.runOnlyPendingTimers();
-            vi.useRealTimers();
-            globalThis.fetch = originalFetch;
-        }
-    });
-});

web/tests/urlUtils.test.js

@@ -106,6 +106,11 @@ describe("buildWorkerWebSocketUrl", () => {
             "wss://worker.example.com/distributed/worker_ws"
         );
     });
+
+    it("always appends /distributed/worker_ws", () => {
+        const url = buildWorkerWebSocketUrl("http://worker.local:8188");
+        expect(url.endsWith("/distributed/worker_ws")).toBe(true);
+    });
 });
 
 

web/tests/workerLifecycle.test.js

@@ -3,7 +3,7 @@ import { afterEach, beforeEach, describe, expect, it } from "vitest";
 import { getWorkerUrl } from "../workerLifecycle.js";
 
 
-describe("workerLifecycle URL wiring", () => {
+describe("workerLifecycle URL construction", () => {
     let originalWindow;
 
     beforeEach(() => {
@@ -22,9 +22,31 @@ describe("workerLifecycle URL wiring", () => {
         globalThis.window = originalWindow;
     });
 
-    // URL variants are covered in urlUtils.test.js; keep the wrapper wiring here.
-    it("forwards the worker and endpoint using window.location", () => {
+    it("builds local worker URL with explicit local port", () => {
         const worker = { id: "w1", port: 8189, type: "local" };
         expect(getWorkerUrl({}, worker, "/prompt")).toBe("http://127.0.0.1:8189/prompt");
     });
+
+    it("builds remote worker URL with host:port", () => {
+        const worker = { id: "w2", host: "worker.example.com", port: 9000, type: "remote" };
+        expect(getWorkerUrl({}, worker, "/prompt")).toBe("http://worker.example.com:9000/prompt");
+    });
+
+    it("builds cloud worker URL as https", () => {
+        const worker = { id: "w3", host: "cloud.example.com", port: 443, type: "cloud" };
+        expect(getWorkerUrl({}, worker, "/prompt")).toBe("https://cloud.example.com/prompt");
+    });
+
+    it("rewrites runpod proxy hostname for local worker ports", () => {
+        globalThis.window = {
+            location: {
+                hostname: "podabc.proxy.runpod.net",
+                protocol: "https:",
+                port: "",
+                origin: "https://podabc.proxy.runpod.net",
+            },
+        };
+        const worker = { id: "w4", port: 8189, type: "local" };
+        expect(getWorkerUrl({}, worker, "/prompt")).toBe("https://podabc-8189.proxy.runpod.net/prompt");
+    });
 });

web/tests/workerSettingsLayout.test.js

@@ -1,16 +0,0 @@
-import { describe, expect, it } from "vitest";
-import { PULSE_ANIMATION_CSS } from "../constants.js";
-
-describe("worker settings layout", () => {
-    it("lets expanded forms grow instead of clipping the final controls", () => {
-        const expanded = PULSE_ANIMATION_CSS.match(/\.worker-settings\.expanded\s*\{([^}]+)\}/)?.[1];
-        expect(expanded).toBeDefined();
-        expect(expanded).toMatch(/max-height:\s*none\s*;/);
-    });
-
-    it("still hides collapsed settings", () => {
-        const collapsed = PULSE_ANIMATION_CSS.match(/\.worker-settings\s*\{([^}]+)\}/)?.[1];
-        expect(collapsed).toMatch(/max-height:\s*0\s*;/);
-        expect(collapsed).toMatch(/overflow:\s*hidden\s*;/);
-    });
-});

web/tests/workerToggle.test.js

@@ -1,54 +0,0 @@
-import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
-
-vi.mock("../../scripts/app.js", () => ({ app: { registerExtension: vi.fn() } }), { virtual: true });
-vi.mock("../../scripts/api.js", () => ({ api: { clientId: "test-client", queuePrompt: vi.fn() } }), { virtual: true });
-
-import { DistributedExtension } from "../main.js";
-
-function makeExtension(worker, updateWorker) {
-    return {
-        config: { workers: [worker] },
-        state: { updateWorker: vi.fn() },
-        api: { updateWorker },
-        app: { extensionManager: { toast: { add: vi.fn() } } },
-        ui: { updateStatusDot: vi.fn() },
-        log: vi.fn(),
-        _emitWorkersChanged: vi.fn(),
-        panelElement: null,
-    };
-}
-
-describe("worker enable toggle", () => {
-    let originalDocument;
-    beforeEach(() => {
-        originalDocument = globalThis.document;
-        globalThis.document = { getElementById: () => null };
-    });
-    afterEach(() => {
-        globalThis.document = originalDocument;
-    });
-
-    it("reverts the toggle and explains why when the backend rejects the change", async () => {
-        const worker = { id: "w1", type: "serverless", enabled: false };
-        const extension = makeExtension(worker, vi.fn().mockRejectedValue(new Error("Save an API key before enabling this worker")));
-
-        await DistributedExtension.prototype.updateWorkerEnabled.call(extension, "w1", true);
-
-        expect(worker.enabled).toBe(false);
-        expect(extension.state.updateWorker).toHaveBeenLastCalledWith("w1", { enabled: false });
-        expect(extension.app.extensionManager.toast.add).toHaveBeenCalledWith(expect.objectContaining({
-            severity: "error", summary: "Could not enable worker", detail: "Save an API key before enabling this worker",
-        }));
-    });
-
-    it("keeps the new state when the backend accepts it", async () => {
-        const worker = { id: "w2", type: "serverless", enabled: false };
-        const extension = makeExtension(worker, vi.fn().mockResolvedValue({ status: "success" }));
-
-        await DistributedExtension.prototype.updateWorkerEnabled.call(extension, "w2", true);
-
-        expect(worker.enabled).toBe(true);
-        expect(extension.state.updateWorker).toHaveBeenLastCalledWith("w2", { enabled: true });
-        expect(extension.app.extensionManager.toast.add).not.toHaveBeenCalled();
-    });
-});

web/ui.js

@@ -7,14 +7,6 @@ import { createLogModal } from './ui/logModal.js';
 import { launchWorker, stopWorker, updateWorkerControls, viewWorkerLog } from './workerLifecycle.js';
 import { isRemoteWorker } from './workerSettings.js';
 
-const escapeWorkerInfo = (value) => String(value ?? "").replace(/[&<>"']/g, (character) => ({
-    "&": "&amp;",
-    "<": "&lt;",
-    ">": "&gt;",
-    '"': "&quot;",
-    "'": "&#39;",
-}[character]));
-
 const cardConfigs = {
     master: {
         checkbox: { 
@@ -66,12 +58,7 @@ const cardConfigs = {
         infoText: (data, extension) => {
             const isRemote = isRemoteWorker(extension, data);
             const isCloud = data.type === 'cloud';
-            const isServerless = data.type === 'serverless';
-
-            if (isServerless) {
-                return `<span class="dist-worker-info__title">${escapeWorkerInfo(data.name)}</span>`;
-            }
-
+            
             if (isCloud) {
                 // For cloud workers, don't show port (it's always 443)
                 return `<span class="dist-worker-info__title">${data.name}</span><br><span class="dist-worker-info__meta">${data.host}</span>`;
@@ -603,10 +590,7 @@ export class DistributedUI {
             }
             controlsWrapper.appendChild(badge);
         } else if (config.dynamic && data) {
-            if (data.type === 'serverless') {
-                // Match the master/remote footer without exposing direct-worker actions.
-                controlsWrapper.appendChild(this.createInfoBox("Runpod Serverless endpoint"));
-            } else if (isRemote) {
+            if (isRemote) {
                 const isCloud = data.type === 'cloud';
                 const workerTypeText = isCloud ? "Cloud worker" : "Remote worker";
                 const workerTypeBadge = this.createInfoBox(workerTypeText);