robertvoy-ComfyUI-Distributed/utils/async_helpers.py

"""
Async helper utilities for ComfyUI-Distributed.
"""
import asyncio
import contextlib
import threading
import time
import uuid
import execution
import server
from typing import Optional, Any, Coroutine
from .network import get_server_loop

def run_async_in_server_loop(coro: Coroutine, timeout: Optional[float] = None) -> Any:
    """
    Run async coroutine in server's event loop and wait for result.
    
    This is useful when you need to run async code from a synchronous context
    but want to use the server's existing event loop instead of creating a new one.
    
    Args:
        coro: The coroutine to run
        timeout: Optional timeout in seconds
        
    Returns:
        The result of the coroutine
        
    Raises:
        TimeoutError: If the operation times out
        Exception: Any exception raised by the coroutine
    """
    event = threading.Event()
    result = None
    error = None
    
    async def wrapper():
        nonlocal result, error
        try:
            result = await coro
        except Exception as e:
            error = e
        finally:
            event.set()
    
    # Schedule on server's event loop
    loop = get_server_loop()
    asyncio.run_coroutine_threadsafe(wrapper(), loop)
    
    # Wait for completion
    if not event.wait(timeout):
        raise TimeoutError(f"Async operation timed out after {timeout} seconds")
    
    if error:
        raise error
    return result


prompt_server = server.PromptServer.instance


def _summarize_node_errors(node_errors: dict) -> str:
    if not isinstance(node_errors, dict) or not node_errors:
        return ""

    parts = []
    for node_id, entry in node_errors.items():
        if not isinstance(entry, dict):
            continue
        class_type = str(entry.get("class_type") or "UnknownNode")
        for err in entry.get("errors", []):
            if not isinstance(err, dict):
                continue
            message = str(err.get("message") or "validation error")
            details = str(err.get("details") or "").strip()
            parts.append(
                f"{class_type}#{node_id}: {message}{f' ({details})' if details else ''}"
            )
            if len(parts) >= 5:
                return " | ".join(parts)
    return " | ".join(parts)


class QueueInsertionCancelled(RuntimeError):
    """The prompt was valid but its insertion was refused by ``insert_if``."""

    cancelled_insertion = True


class PromptValidationError(RuntimeError):
    """Raised when a prompt fails ComfyUI validation with structured context."""

    def __init__(self, error_payload, node_errors=None):
        payload = error_payload if isinstance(error_payload, dict) else {
            "type": "prompt_validation_failed",
            "message": str(error_payload),
            "details": "",
            "extra_info": {},
        }
        self.validation_error = dict(payload)
        self.node_errors = node_errors if isinstance(node_errors, dict) else {}

        if self.node_errors:
            details = str(self.validation_error.get("details") or "").strip()
            if not details:
                summary = _summarize_node_errors(self.node_errors)
                if summary:
                    self.validation_error["details"] = summary

        merged = dict(self.validation_error)
        if self.node_errors:
            merged["node_errors"] = self.node_errors
        super().__init__(f"Invalid prompt: {merged}")


def validate_execution_targets(prompt, targets):
    """Reject malformed explicit roots instead of falling back to all outputs."""
    if (not isinstance(targets, list) or not targets
            or any(not isinstance(node_id, str) or not node_id or node_id not in prompt
                   for node_id in targets)
            or len(set(targets)) != len(targets)):
        raise PromptValidationError({
            "type": "invalid_partial_execution_targets",
            "message": "partial_execution_targets must be a non-empty list of unique existing node IDs",
            "details": "", "extra_info": {},
        })
    return list(targets)


async def queue_prompt_payload(
    prompt_obj,
    workflow_meta=None,
    client_id=None,
    include_queue_metadata=False,
    *,
    partial_execution_targets=None,
    prompt_id=None,
    insert_if=None,
):
    """Validate and queue a prompt, preserving any explicitly requested roots.

    ``prompt_id`` lets a caller re-queue a prompt under an id it already
    handed out (the load-balance dispatcher swaps a waiting prompt for its
    delegate form without changing the id clients track).

    ``insert_if`` is checked under the queue mutex right before insertion,
    after validation has yielded to the event loop; when it returns false
    the prompt is not queued and ``QueueInsertionCancelled`` is raised.
    """
    targets = (validate_execution_targets(prompt_obj, partial_execution_targets)
               if partial_execution_targets is not None else None)
    payload = {"prompt": prompt_obj}
    if targets is not None:
        payload["partial_execution_targets"] = list(targets)
    payload = prompt_server.trigger_on_prompt(payload)
    prompt = payload["prompt"]
    if targets is not None:
        # Hooks may replace the envelope or mutate its target list. The caller's
        # requested roots remain authoritative, but must still exist afterwards.
        validate_execution_targets(prompt, targets)

    prompt_id = str(prompt_id) if prompt_id else str(uuid.uuid4())
    valid = await execution.validate_prompt(prompt_id, prompt, targets)
    if not valid[0]:
        error_payload = valid[1] if len(valid) > 1 else "Prompt outputs failed validation"
        node_errors = valid[3] if len(valid) > 3 else {}
        raise PromptValidationError(error_payload, node_errors)

    if targets is not None:
        node_errors = valid[3] if len(valid) > 3 else {}
        target_errors = {
            node_id: entry for node_id, entry in node_errors.items()
            if node_id in targets or (isinstance(entry, dict)
                                      and set(entry.get("dependent_outputs", [])) & set(targets))
        }
        if set(valid[2]) != set(targets) or target_errors:
            raise PromptValidationError({
                "type": "partial_execution_targets_failed_validation",
                "message": "Not all requested execution targets passed validation",
                "details": "Refusing to queue a subset or additional output roots",
                "extra_info": {"requested_targets": targets},
            }, node_errors)

    extra_data = {"create_time": int(time.time() * 1000)}
    if workflow_meta:
        extra_data.setdefault("extra_pnginfo", {})["workflow"] = workflow_meta
    if client_id:
        extra_data["client_id"] = client_id

    sensitive = {}
    for key in getattr(execution, "SENSITIVE_EXTRA_DATA_KEYS", []):
        if key in extra_data:
            sensitive[key] = extra_data.pop(key)

    prompt_queue = prompt_server.prompt_queue
    mutex = getattr(prompt_queue, "mutex", None)
    with (mutex if mutex is not None else contextlib.nullcontext()):
        if insert_if is not None and not insert_if():
            raise QueueInsertionCancelled(f"Prompt {prompt_id} was cancelled before it could be queued")
        number = getattr(prompt_server, "number", 0)
        prompt_server.number = number + 1
        prompt_queue_item = (number, prompt_id, prompt, extra_data, valid[2], sensitive)
        prompt_queue.put(prompt_queue_item)

    if include_queue_metadata:
        return {
            "prompt_id": prompt_id,
            "number": number,
            "node_errors": {},
        }

    return prompt_id