Sebastian Monroy
66c66f045f
Nilor-Nodes worker websocket listener was failing with: TypeError: BaseEventLoop.create_connection() got an unexpected keyword argument 'read_limit' On Python 3.12 with websockets>=15, passing read_limit to websockets.connect() ends up forwarding it to asyncio’s loop.create_connection(), which doesn’t accept that kwarg. This change removes the read_limit argument from ComfyUILocalClient.ws_connect() while keeping max_size=None and existing ping settings, restoring websocket connectivity and preventing reconnect exhaustion.
911 lines
32 KiB
Python
911 lines
32 KiB
Python
"""
|
||
Thin, typed client surface for accessing ComfyUI HTTP endpoints and the websocket.
|
||
|
||
This module defines the public protocol, DTOs, and exceptions that callers and
|
||
tests depend on. Implementations are intentionally minimal at this stage; network
|
||
behavior will be added in subsequent commits.
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import asyncio
|
||
import json
|
||
import random
|
||
from dataclasses import dataclass
|
||
from typing import Any, AsyncIterator, Dict, Optional, Protocol, TypedDict, Tuple
|
||
|
||
from urllib.parse import quote, urlparse
|
||
|
||
|
||
__all__ = [
|
||
"ComfyUIClientProtocol",
|
||
"ComfyUILocalClient",
|
||
"SystemStats",
|
||
"WsEvent",
|
||
"ComfyUIClientError",
|
||
"ComfyUIClientTimeout",
|
||
"ComfyUIClientWsClosed",
|
||
]
|
||
|
||
|
||
class WsEvent(TypedDict, total=False):
|
||
"""Typed view of websocket events emitted by ComfyUI.
|
||
|
||
Fields:
|
||
- type: Event type string (e.g., "status", "progress", "executed").
|
||
- data: Opaque payload; commonly includes keys like "prompt_id", "node", etc.
|
||
"""
|
||
|
||
type: str
|
||
data: Dict[str, Any]
|
||
|
||
|
||
@dataclass
|
||
class SystemStats:
|
||
"""Subset of system statistics reported by ComfyUI `/system_stats`.
|
||
|
||
Known fields are optional; unknown fields should be ignored by parsers. When
|
||
present, RAM-related metrics may also appear (e.g., `ram_total`, `ram_free`).
|
||
|
||
Attributes:
|
||
vram_total: Total VRAM (bytes) when reported.
|
||
vram_free: Free VRAM reported by the backend, if available.
|
||
torch_vram_free: Free VRAM according to torch, if available.
|
||
ram_total: Total system RAM (bytes) when reported.
|
||
ram_free: Free system RAM (bytes) when reported.
|
||
"""
|
||
|
||
vram_total: Optional[float] = None
|
||
vram_free: Optional[float] = None
|
||
torch_vram_free: Optional[float] = None
|
||
ram_total: Optional[float] = None
|
||
ram_free: Optional[float] = None
|
||
|
||
|
||
class ComfyUIClientError(Exception):
|
||
"""Base error for all ComfyUI client failures.
|
||
|
||
Args:
|
||
message: Human-friendly error message.
|
||
route: Route path (e.g., "/prompt").
|
||
method: HTTP method (e.g., "GET", "POST").
|
||
status: Optional HTTP status code or websocket close code.
|
||
code: Optional machine-readable error code (e.g., "timeout").
|
||
body_snippet: Optional diagnostic snippet from a response payload.
|
||
"""
|
||
|
||
def __init__(
|
||
self,
|
||
message: str,
|
||
*,
|
||
route: Optional[str] = None,
|
||
method: Optional[str] = None,
|
||
status: Optional[int] = None,
|
||
code: Optional[str] = None,
|
||
body_snippet: Optional[str] = None,
|
||
) -> None:
|
||
super().__init__(message)
|
||
self.route: Optional[str] = route
|
||
self.method: Optional[str] = method
|
||
self.status: Optional[int] = status
|
||
self.code: Optional[str] = code
|
||
self.body_snippet: Optional[str] = body_snippet
|
||
|
||
|
||
class ComfyUIClientTimeout(ComfyUIClientError):
|
||
"""Raised when an operation exceeds its allowed timeout."""
|
||
|
||
|
||
class ComfyUIClientWsClosed(ComfyUIClientError):
|
||
"""Raised when the websocket is closed or cannot be maintained."""
|
||
|
||
|
||
class ComfyUIClientProtocol(Protocol):
|
||
"""Protocol for ComfyUI clients.
|
||
|
||
Callers and tests should depend on this interface rather than a concrete
|
||
implementation. Methods are asynchronous and may raise subclasses of
|
||
`ComfyUIClientError`.
|
||
"""
|
||
|
||
async def submit_prompt(self, payload: Dict[str, Any]) -> str:
|
||
"""Submit a prompt to ComfyUI and return the resulting `prompt_id`.
|
||
|
||
Args:
|
||
payload: JSON-serializable payload for the `/prompt` endpoint.
|
||
|
||
Returns:
|
||
The non-empty `prompt_id` string returned by the server.
|
||
"""
|
||
|
||
async def get_system_stats(self) -> SystemStats:
|
||
"""Fetch system statistics from `/system_stats`."""
|
||
|
||
async def free(
|
||
self, *, free_memory: bool = False, unload_models: bool = False
|
||
) -> None:
|
||
"""Invoke `/free` with the provided flags."""
|
||
|
||
async def ws_connect(self, client_id: str) -> AsyncIterator[WsEvent]:
|
||
"""Connect to the websocket (`/ws?clientId=...`) and yield parsed events."""
|
||
|
||
async def probe(self) -> None:
|
||
"""Lightweight health probe for connectivity/parseability.
|
||
|
||
Executes a GET `/system_stats` with a short timeout and no retries. Raises
|
||
`ComfyUIClientError` subclasses on failure; returns `None` on success.
|
||
"""
|
||
|
||
async def supports_hygiene(self) -> bool:
|
||
"""Return True if both `/system_stats` and `/free` are supported.
|
||
|
||
Performs a one-time capability probe; caches results for the session.
|
||
"""
|
||
|
||
|
||
class ComfyUILocalClient(ComfyUIClientProtocol):
|
||
"""Local HTTP/WebSocket client for a running ComfyUI instance.
|
||
|
||
This class provides the concrete implementation for the protocol. At this
|
||
stage it only declares the interface and stores constructor parameters; the
|
||
network behavior will be implemented in subsequent commits.
|
||
|
||
Args:
|
||
base_url: Base HTTP URL for ComfyUI endpoints (e.g., `/prompt`).
|
||
ws_url: Base WebSocket URL (e.g., `/ws`).
|
||
session: Optional externally-managed aiohttp session for reuse.
|
||
logger: Optional logger compatible with the worker's logging API.
|
||
timeout: Default timeout in seconds for HTTP operations.
|
||
"""
|
||
|
||
def __init__(
|
||
self,
|
||
base_url: str,
|
||
ws_url: str,
|
||
session: Optional["aiohttp.ClientSession"] = None,
|
||
logger: Optional[Any] = None,
|
||
timeout: float = 30.0,
|
||
) -> None:
|
||
self._base_url: str = base_url
|
||
self._ws_url: str = ws_url
|
||
self._session = session # type: ignore[assignment]
|
||
self._logger = logger
|
||
self._timeout: float = float(timeout)
|
||
self._owned_session: Optional[aiohttp.ClientSession] = None
|
||
# Backoff defaults per plan (commit 3)
|
||
self._retry_base_seconds: float = 0.25
|
||
self._retry_multiplier: float = 2.0
|
||
self._retry_jitter_seconds: float = 0.25
|
||
self._retry_max_sleep_seconds: float = 4.0
|
||
self._retry_max_attempts: int = 3
|
||
# WebSocket reconnect policy (commit 4)
|
||
self._ws_max_reconnect_attempts: int = 5
|
||
self._ws_max_total_backoff_seconds: float = 30.0
|
||
# Capability probe cache (None = unknown, True/False = probed)
|
||
self._supports_system_stats: Optional[bool] = None
|
||
self._supports_free: Optional[bool] = None
|
||
self._capability_warning_emitted: bool = False
|
||
|
||
# Lifecycle methods may be implemented later; for now they act as no-ops.
|
||
async def __aenter__(self) -> "ComfyUILocalClient":
|
||
"""Enter async context; create an internal session when none provided."""
|
||
if self._session is None and self._owned_session is None:
|
||
self._owned_session = aiohttp.ClientSession()
|
||
return self
|
||
|
||
async def __aexit__(self, exc_type, exc, tb) -> None: # type: ignore[override]
|
||
"""Exit async context; close internal session if owned by this client."""
|
||
if self._owned_session is not None:
|
||
try:
|
||
await self._owned_session.close()
|
||
finally:
|
||
self._owned_session = None
|
||
return None
|
||
|
||
# Protocol methods — to be implemented in subsequent commits.
|
||
async def submit_prompt(self, payload: Dict[str, Any]) -> str: # type: ignore[override]
|
||
route = "/prompt"
|
||
method = "POST"
|
||
url = self._join_http(route)
|
||
try:
|
||
data = await self._http_request_json(
|
||
method,
|
||
url,
|
||
json_payload=payload,
|
||
timeout_s=self._timeout,
|
||
retry_idempotent=False,
|
||
)
|
||
except asyncio.TimeoutError as e:
|
||
raise ComfyUIClientTimeout(
|
||
f"Timeout while calling {method} {route}",
|
||
route=route,
|
||
method=method,
|
||
code="timeout",
|
||
) from e
|
||
except _MappedHttpError as e:
|
||
raise e.to_public_error(route=route, method=method)
|
||
except _MappedConnError as e:
|
||
raise e.to_public_error(route=route, method=method)
|
||
|
||
prompt_id = data.get("prompt_id") if isinstance(data, dict) else None
|
||
if not isinstance(prompt_id, str) or not prompt_id.strip():
|
||
snippet = _safe_preview(data)
|
||
raise ComfyUIClientError(
|
||
"Invalid response payload: missing non-empty prompt_id",
|
||
route=route,
|
||
method=method,
|
||
code="bad_json",
|
||
body_snippet=snippet,
|
||
)
|
||
return prompt_id
|
||
|
||
async def get_system_stats(self) -> SystemStats: # type: ignore[override]
|
||
route = "/system_stats"
|
||
method = "GET"
|
||
url = self._join_http(route)
|
||
try:
|
||
data = await self._http_request_json(
|
||
method,
|
||
url,
|
||
json_payload=None,
|
||
timeout_s=self._timeout,
|
||
retry_idempotent=True,
|
||
)
|
||
except asyncio.TimeoutError as e:
|
||
raise ComfyUIClientTimeout(
|
||
f"Timeout while calling {method} {route}",
|
||
route=route,
|
||
method=method,
|
||
code="timeout",
|
||
) from e
|
||
except _MappedHttpError as e:
|
||
raise e.to_public_error(route=route, method=method)
|
||
except _MappedConnError as e:
|
||
raise e.to_public_error(route=route, method=method)
|
||
|
||
# Parse known fields, tolerate missing/unknown; try alternate shapes; fallback to torch
|
||
vram_total = None
|
||
vram_free = None
|
||
torch_vram_free = None
|
||
ram_total = None
|
||
ram_free = None
|
||
|
||
if isinstance(data, dict):
|
||
vram_total = _coerce_optional_float(data.get("vram_total"))
|
||
vram_free = _coerce_optional_float(data.get("vram_free"))
|
||
torch_vram_free = _coerce_optional_float(data.get("torch_vram_free"))
|
||
ram_total = _coerce_optional_float(data.get("ram_total"))
|
||
ram_free = _coerce_optional_float(data.get("ram_free"))
|
||
|
||
# Common alternates
|
||
if vram_total is None:
|
||
vram_total = _coerce_optional_float(data.get("total_vram"))
|
||
if vram_free is None:
|
||
vram_free = _coerce_optional_float(data.get("free_vram"))
|
||
|
||
vram_obj = data.get("vram") if isinstance(data.get("vram"), dict) else None
|
||
if vram_obj:
|
||
if vram_total is None:
|
||
vram_total = _coerce_optional_float(vram_obj.get("total"))
|
||
if vram_free is None:
|
||
vram_free = _coerce_optional_float(vram_obj.get("free"))
|
||
|
||
ram_obj = data.get("ram") if isinstance(data.get("ram"), dict) else None
|
||
if ram_obj:
|
||
if ram_total is None:
|
||
ram_total = _coerce_optional_float(ram_obj.get("total"))
|
||
if ram_free is None:
|
||
ram_free = _coerce_optional_float(ram_obj.get("free"))
|
||
|
||
# devices[0] fallback (common in mock/alt servers)
|
||
devices = (
|
||
data.get("devices") if isinstance(data.get("devices"), list) else None
|
||
)
|
||
if devices and len(devices) > 0 and isinstance(devices[0], dict):
|
||
dev0 = devices[0]
|
||
if vram_total is None:
|
||
vram_total = _coerce_optional_float(dev0.get("vram_total"))
|
||
if vram_free is None:
|
||
vram_free = _coerce_optional_float(dev0.get("vram_free"))
|
||
# Note: we rely solely on server-reported values; no local torch fallback
|
||
|
||
# Optional debug logging of reported stats
|
||
if self._logger:
|
||
try:
|
||
self._logger.debug(
|
||
"ℹ️\u2009 Nilor-Nodes (comfyui_client): /system_stats: vram_total=%s vram_free=%s torch_vram_free=%s ram_total=%s ram_free=%s",
|
||
vram_total,
|
||
vram_free,
|
||
torch_vram_free,
|
||
ram_total,
|
||
ram_free,
|
||
)
|
||
except Exception:
|
||
pass
|
||
return SystemStats(
|
||
vram_total=vram_total,
|
||
vram_free=vram_free,
|
||
torch_vram_free=torch_vram_free,
|
||
ram_total=ram_total,
|
||
ram_free=ram_free,
|
||
)
|
||
|
||
async def free(self, *, free_memory: bool = False, unload_models: bool = False) -> None: # type: ignore[override]
|
||
route = "/free"
|
||
method = "POST"
|
||
url = self._join_http(route)
|
||
body = {"free_memory": bool(free_memory), "unload_models": bool(unload_models)}
|
||
try:
|
||
await self._http_request_json(
|
||
method,
|
||
url,
|
||
json_payload=body,
|
||
timeout_s=self._timeout,
|
||
retry_idempotent=True, # idempotent when flags identical
|
||
)
|
||
except asyncio.TimeoutError as e:
|
||
raise ComfyUIClientTimeout(
|
||
f"Timeout while calling {method} {route}",
|
||
route=route,
|
||
method=method,
|
||
code="timeout",
|
||
) from e
|
||
except _MappedHttpError as e:
|
||
raise e.to_public_error(route=route, method=method)
|
||
except _MappedConnError as e:
|
||
raise e.to_public_error(route=route, method=method)
|
||
|
||
async def ws_connect(self, client_id: str) -> AsyncIterator[WsEvent]: # type: ignore[override]
|
||
url = _ws_url(self._ws_url, client_id)
|
||
attempts = 0
|
||
total_backoff = 0.0
|
||
base = self._retry_base_seconds
|
||
multiplier = self._retry_multiplier
|
||
jitter = self._retry_jitter_seconds
|
||
max_sleep = self._retry_max_sleep_seconds
|
||
|
||
while True:
|
||
try:
|
||
# Allow large frames and set ping/pong defaults
|
||
async with websockets.connect(
|
||
url,
|
||
max_size=None,
|
||
max_queue=4,
|
||
ping_interval=20,
|
||
ping_timeout=20,
|
||
) as websocket:
|
||
# On successful connect, reset counters
|
||
attempts = 0
|
||
total_backoff = 0.0
|
||
if self._logger:
|
||
try:
|
||
self._logger.debug(
|
||
f"✅ Nilor-Nodes (comfyui_client): connected to websocket {url}"
|
||
)
|
||
except Exception:
|
||
pass
|
||
|
||
while True:
|
||
message = await websocket.recv()
|
||
if isinstance(message, bytes):
|
||
try:
|
||
message = message.decode("utf-8", errors="replace")
|
||
except Exception:
|
||
yield WsEvent(type="binary", data={"length": len(message)}) # type: ignore[call-arg]
|
||
continue
|
||
try:
|
||
payload = json.loads(message)
|
||
except Exception:
|
||
yield WsEvent(type="text", data={"message": message}) # type: ignore[call-arg]
|
||
continue
|
||
|
||
if isinstance(payload, dict):
|
||
event_type = str(payload.get("type", "event"))
|
||
event_data = payload.get("data")
|
||
if not isinstance(event_data, dict):
|
||
event_data = {"raw": payload}
|
||
yield WsEvent(type=event_type, data=event_data) # type: ignore[call-arg]
|
||
else:
|
||
yield WsEvent(type="event", data={"raw": payload}) # type: ignore[call-arg]
|
||
|
||
except asyncio.CancelledError:
|
||
# Allow clean shutdown by propagating cancellation
|
||
raise
|
||
except ws_exc.ConnectionClosedOK as e:
|
||
raise ComfyUIClientWsClosed(
|
||
"WebSocket closed normally",
|
||
route="/ws",
|
||
method="GET",
|
||
status=getattr(e, "code", 1000),
|
||
code="ws_closed",
|
||
body_snippet=str(getattr(e, "reason", ""))[:256],
|
||
) from e
|
||
except ws_exc.ConnectionClosedError as e:
|
||
# Abnormal close: attempt bounded reconnect
|
||
attempts += 1
|
||
if attempts > self._ws_max_reconnect_attempts:
|
||
raise ComfyUIClientError(
|
||
"WebSocket reconnect attempts exhausted",
|
||
route="/ws",
|
||
method="GET",
|
||
code="ws_closed",
|
||
body_snippet=str(getattr(e, "reason", ""))[:256],
|
||
) from e
|
||
delay = _next_backoff(attempts - 1, base, multiplier, jitter, max_sleep)
|
||
total_backoff += delay
|
||
if total_backoff > self._ws_max_total_backoff_seconds:
|
||
raise ComfyUIClientError(
|
||
"WebSocket reconnect backoff budget exhausted",
|
||
route="/ws",
|
||
method="GET",
|
||
code="ws_closed",
|
||
body_snippet=str(getattr(e, "reason", ""))[:256],
|
||
) from e
|
||
await asyncio.sleep(delay)
|
||
continue
|
||
except ws_exc.InvalidStatus as e:
|
||
raise ComfyUIClientError(
|
||
"WebSocket handshake failed",
|
||
route="/ws",
|
||
method="GET",
|
||
code="ws_handshake",
|
||
) from e
|
||
except ws_exc.InvalidURI as e:
|
||
raise ComfyUIClientError(
|
||
"Invalid WebSocket URI",
|
||
route="/ws",
|
||
method="GET",
|
||
code="invalid_ws_uri",
|
||
) from e
|
||
except Exception as e:
|
||
# Treat as connection error; bounded reconnect
|
||
attempts += 1
|
||
if attempts > self._ws_max_reconnect_attempts:
|
||
raise ComfyUIClientError(
|
||
"WebSocket reconnect attempts exhausted",
|
||
route="/ws",
|
||
method="GET",
|
||
code="connection_error",
|
||
) from e
|
||
delay = _next_backoff(attempts - 1, base, multiplier, jitter, max_sleep)
|
||
total_backoff += delay
|
||
if total_backoff > self._ws_max_total_backoff_seconds:
|
||
raise ComfyUIClientError(
|
||
"WebSocket reconnect backoff budget exhausted",
|
||
route="/ws",
|
||
method="GET",
|
||
code="connection_error",
|
||
) from e
|
||
await asyncio.sleep(delay)
|
||
continue
|
||
|
||
async def probe(self) -> None: # type: ignore[override]
|
||
route = "/system_stats"
|
||
method = "GET"
|
||
url = self._join_http(route)
|
||
short_timeout = min(self._timeout, 3.0)
|
||
try:
|
||
# No retries: retry_idempotent=False
|
||
await self._http_request_json(
|
||
method,
|
||
url,
|
||
json_payload=None,
|
||
timeout_s=short_timeout,
|
||
retry_idempotent=False,
|
||
)
|
||
except asyncio.TimeoutError as e:
|
||
raise ComfyUIClientTimeout(
|
||
f"Timeout while calling {method} {route}",
|
||
route=route,
|
||
method=method,
|
||
code="timeout",
|
||
) from e
|
||
except _MappedHttpError as e:
|
||
raise e.to_public_error(route=route, method=method)
|
||
except _MappedConnError as e:
|
||
raise e.to_public_error(route=route, method=method)
|
||
|
||
async def supports_hygiene(self) -> bool: # type: ignore[override]
|
||
# If both probed, return cached decision
|
||
if self._supports_system_stats is not None and self._supports_free is not None:
|
||
return bool(self._supports_system_stats and self._supports_free)
|
||
|
||
await self._probe_capabilities_once()
|
||
supported = bool(
|
||
(self._supports_system_stats is True) and (self._supports_free is True)
|
||
)
|
||
if not supported and not self._capability_warning_emitted and self._logger:
|
||
try:
|
||
self._logger.warning(
|
||
"⚠️\u2009 Nilor-Nodes (comfyui_client): hygiene disabled — missing /system_stats or /free support"
|
||
)
|
||
except Exception:
|
||
pass
|
||
self._capability_warning_emitted = True
|
||
return supported
|
||
|
||
async def _probe_capabilities_once(self) -> None:
|
||
"""Probe `/system_stats` and `/free` capabilities once and cache results.
|
||
|
||
Only marks capabilities as False on definitive 404/405 responses. Transient
|
||
failures leave the capability as None so a future call may retry.
|
||
"""
|
||
short_timeout = min(self._timeout, 3.0)
|
||
|
||
# Probe /system_stats support
|
||
if self._supports_system_stats is None:
|
||
route = "/system_stats"
|
||
url = self._join_http(route)
|
||
try:
|
||
await self._http_request_json(
|
||
"GET",
|
||
url,
|
||
json_payload=None,
|
||
timeout_s=short_timeout,
|
||
retry_idempotent=False,
|
||
)
|
||
self._supports_system_stats = True
|
||
except ComfyUIClientError as e:
|
||
if getattr(e, "status", None) in (404, 405):
|
||
self._supports_system_stats = False
|
||
|
||
# Probe /free support (no-op body)
|
||
if self._supports_free is None:
|
||
route = "/free"
|
||
url = self._join_http(route)
|
||
try:
|
||
await self._http_request_json(
|
||
"POST",
|
||
url,
|
||
json_payload={"free_memory": False, "unload_models": False},
|
||
timeout_s=short_timeout,
|
||
retry_idempotent=False,
|
||
)
|
||
self._supports_free = True
|
||
except ComfyUIClientError as e:
|
||
if getattr(e, "status", None) in (404, 405):
|
||
self._supports_free = False
|
||
|
||
|
||
# Runtime dependency; imported here to avoid issues if module is scanned without execution
|
||
import aiohttp # type: ignore
|
||
import websockets # type: ignore
|
||
from websockets import exceptions as ws_exc # type: ignore
|
||
|
||
|
||
# ---- Internal helpers (HTTP) ----
|
||
|
||
|
||
def _coerce_optional_float(value: Any) -> Optional[float]:
|
||
try:
|
||
if value is None:
|
||
return None
|
||
return float(value)
|
||
except Exception:
|
||
return None
|
||
|
||
|
||
def _is_transient_status(status: int) -> bool:
|
||
return status == 429 or 500 <= status <= 599
|
||
|
||
|
||
def _safe_preview(data: Any, limit: int = 512) -> str:
|
||
try:
|
||
text = json.dumps(data, ensure_ascii=False)
|
||
except Exception:
|
||
text = str(data)
|
||
if len(text) > limit:
|
||
return text[:limit] + "…"
|
||
return text
|
||
|
||
|
||
class _MappedHttpError(Exception):
|
||
def __init__(self, *, status: Optional[int], body_snippet: Optional[str]) -> None:
|
||
self.status = status
|
||
self.body_snippet = body_snippet
|
||
|
||
def to_public_error(self, *, route: str, method: str) -> ComfyUIClientError:
|
||
return ComfyUIClientError(
|
||
f"🛑\u2009 Nilor-Nodes (comfyui_client): HTTP error while calling {method} {route}",
|
||
route=route,
|
||
method=method,
|
||
status=self.status,
|
||
code="http_error",
|
||
body_snippet=self.body_snippet,
|
||
)
|
||
|
||
|
||
class _MappedConnError(Exception):
|
||
def __init__(self, *, code: str = "connection_error") -> None:
|
||
self.code = code
|
||
|
||
def to_public_error(self, *, route: str, method: str) -> ComfyUIClientError:
|
||
return ComfyUIClientError(
|
||
f"🛑\u2009 Nilor-Nodes (comfyui_client): Connection error while calling {method} {route}",
|
||
route=route,
|
||
method=method,
|
||
code=self.code,
|
||
)
|
||
|
||
|
||
async def _read_limited_text(resp: aiohttp.ClientResponse, limit: int = 512) -> str:
|
||
try:
|
||
raw = await resp.read()
|
||
# Truncate at byte level then decode safely
|
||
raw = raw[:limit]
|
||
return raw.decode("utf-8", errors="replace")
|
||
except Exception:
|
||
return ""
|
||
|
||
|
||
def _with_jitter(seconds: float, jitter: float) -> float:
|
||
if jitter <= 0:
|
||
return seconds
|
||
return max(0.0, seconds + random.uniform(-jitter, jitter))
|
||
|
||
|
||
def _next_backoff(
|
||
attempt_index: int,
|
||
base: float,
|
||
multiplier: float,
|
||
jitter: float,
|
||
max_sleep: float,
|
||
) -> float:
|
||
# attempt_index is 0-based
|
||
delay = base * (multiplier**attempt_index)
|
||
delay = min(delay, max_sleep)
|
||
return _with_jitter(delay, jitter)
|
||
|
||
|
||
def _should_retry(
|
||
*,
|
||
retry_idempotent: bool,
|
||
exc: Optional[BaseException] = None,
|
||
status: Optional[int] = None,
|
||
) -> bool:
|
||
if not retry_idempotent:
|
||
return False
|
||
if isinstance(exc, (asyncio.TimeoutError, aiohttp.ClientConnectionError)):
|
||
return True
|
||
if status is not None and _is_transient_status(status):
|
||
return True
|
||
return False
|
||
|
||
|
||
def _finalize_attempts(
|
||
method: str,
|
||
url: str,
|
||
*,
|
||
last_exc: Optional[BaseException],
|
||
last_status: Optional[int],
|
||
last_body_snippet: Optional[str],
|
||
) -> BaseException:
|
||
if isinstance(last_exc, asyncio.TimeoutError):
|
||
return ComfyUIClientTimeout(
|
||
f"🛑\u2009 Nilor-Nodes (comfyui_client): Timeout while calling {method} {url}",
|
||
route=_route_from_url(url),
|
||
method=method,
|
||
code="timeout",
|
||
)
|
||
if isinstance(last_exc, aiohttp.ClientConnectionError):
|
||
return _MappedConnError().to_public_error(
|
||
route=_route_from_url(url), method=method
|
||
)
|
||
# Otherwise treat as HTTP error
|
||
return _MappedHttpError(
|
||
status=last_status, body_snippet=last_body_snippet
|
||
).to_public_error(route=_route_from_url(url), method=method)
|
||
|
||
|
||
def _route_from_url(url: str) -> str:
|
||
try:
|
||
parsed = urlparse(url)
|
||
return parsed.path or "/"
|
||
except Exception:
|
||
return url
|
||
|
||
|
||
class _TempSession:
|
||
"""Context manager that yields an aiohttp session, reusing if provided."""
|
||
|
||
def __init__(self, session: Optional[aiohttp.ClientSession]):
|
||
self._provided = session
|
||
self._owned: Optional[aiohttp.ClientSession] = None
|
||
|
||
async def __aenter__(self) -> aiohttp.ClientSession:
|
||
if self._provided is not None:
|
||
return self._provided
|
||
self._owned = aiohttp.ClientSession()
|
||
return self._owned
|
||
|
||
async def __aexit__(self, exc_type, exc, tb) -> None: # type: ignore[override]
|
||
if self._owned is not None:
|
||
await self._owned.close()
|
||
|
||
|
||
async def _json_or_text(resp: aiohttp.ClientResponse) -> Any:
|
||
ctype = resp.headers.get("Content-Type", "").lower()
|
||
text = await _read_limited_text(resp) # limited read to use for both cases
|
||
if "json" in ctype:
|
||
try:
|
||
return json.loads(text)
|
||
except Exception:
|
||
# Fallthrough to treat as bad JSON
|
||
raise _MappedHttpError(status=resp.status, body_snippet=text)
|
||
# Not JSON; return raw text
|
||
return text
|
||
|
||
|
||
async def _raise_for_status_with_snippet(resp: aiohttp.ClientResponse) -> None:
|
||
if 200 <= resp.status <= 299:
|
||
return
|
||
snippet = await _read_limited_text(resp)
|
||
raise _MappedHttpError(status=resp.status, body_snippet=snippet)
|
||
|
||
|
||
async def _request_once(
|
||
method: str,
|
||
url: str,
|
||
*,
|
||
session: aiohttp.ClientSession,
|
||
json_payload: Optional[Dict[str, Any]],
|
||
timeout_s: float,
|
||
) -> Tuple[Optional[int], Optional[str], Any]:
|
||
timeout = aiohttp.ClientTimeout(total=timeout_s)
|
||
try:
|
||
async with session.request(
|
||
method, url, json=json_payload, timeout=timeout
|
||
) as resp:
|
||
status = resp.status
|
||
await _raise_for_status_with_snippet(resp)
|
||
# Success path: attempt to parse JSON body; if not JSON, return text
|
||
try:
|
||
data = await resp.json(content_type=None)
|
||
except aiohttp.ContentTypeError:
|
||
# Not JSON; use limited text
|
||
data = await _read_limited_text(resp)
|
||
return status, None, data
|
||
except asyncio.TimeoutError:
|
||
raise
|
||
except aiohttp.ClientConnectionError as e:
|
||
raise e
|
||
except aiohttp.ClientPayloadError as e:
|
||
# Map as payload error
|
||
raise _MappedHttpError(status=None, body_snippet=str(e))
|
||
|
||
|
||
async def _http_request_core(
|
||
method: str,
|
||
url: str,
|
||
*,
|
||
session: aiohttp.ClientSession,
|
||
json_payload: Optional[Dict[str, Any]],
|
||
timeout_s: float,
|
||
retry_idempotent: bool,
|
||
base: float,
|
||
multiplier: float,
|
||
jitter: float,
|
||
max_sleep: float,
|
||
max_attempts: int,
|
||
) -> Any:
|
||
last_exc: Optional[BaseException] = None
|
||
last_status: Optional[int] = None
|
||
last_body: Optional[str] = None
|
||
|
||
attempts = max(1, int(max_attempts))
|
||
for attempt in range(attempts):
|
||
try:
|
||
status, body_snippet, data = await _request_once(
|
||
method,
|
||
url,
|
||
session=session,
|
||
json_payload=json_payload,
|
||
timeout_s=timeout_s,
|
||
)
|
||
return data
|
||
except asyncio.TimeoutError as e:
|
||
last_exc = e
|
||
if attempt < attempts - 1 and _should_retry(
|
||
retry_idempotent=retry_idempotent, exc=e
|
||
):
|
||
await asyncio.sleep(
|
||
_next_backoff(attempt, base, multiplier, jitter, max_sleep)
|
||
)
|
||
continue
|
||
break
|
||
except aiohttp.ClientConnectionError as e:
|
||
last_exc = e
|
||
if attempt < attempts - 1 and _should_retry(
|
||
retry_idempotent=retry_idempotent, exc=e
|
||
):
|
||
await asyncio.sleep(
|
||
_next_backoff(attempt, base, multiplier, jitter, max_sleep)
|
||
)
|
||
continue
|
||
break
|
||
except _MappedHttpError as e:
|
||
last_exc = None
|
||
last_status = e.status
|
||
last_body = e.body_snippet
|
||
if attempt < attempts - 1 and _should_retry(
|
||
retry_idempotent=retry_idempotent, status=e.status
|
||
):
|
||
await asyncio.sleep(
|
||
_next_backoff(attempt, base, multiplier, jitter, max_sleep)
|
||
)
|
||
continue
|
||
break
|
||
|
||
raise _finalize_attempts(
|
||
method,
|
||
url,
|
||
last_exc=last_exc,
|
||
last_status=last_status,
|
||
last_body_snippet=last_body,
|
||
)
|
||
|
||
|
||
async def _http_request_json(
|
||
self: "ComfyUILocalClient",
|
||
method: str,
|
||
url: str,
|
||
*,
|
||
json_payload: Optional[Dict[str, Any]],
|
||
timeout_s: float,
|
||
retry_idempotent: bool,
|
||
) -> Any:
|
||
async with _TempSession(self._session or self._owned_session) as session:
|
||
return await _http_request_core(
|
||
method,
|
||
url,
|
||
session=session,
|
||
json_payload=json_payload,
|
||
timeout_s=timeout_s,
|
||
retry_idempotent=retry_idempotent,
|
||
base=self._retry_base_seconds,
|
||
multiplier=self._retry_multiplier,
|
||
jitter=self._retry_jitter_seconds,
|
||
max_sleep=self._retry_max_sleep_seconds,
|
||
max_attempts=self._retry_max_attempts,
|
||
)
|
||
|
||
|
||
def _join_http_base(base: str, route: str) -> str:
|
||
if not route:
|
||
return base
|
||
return f"{base.rstrip('/')}{route}"
|
||
|
||
|
||
def _join_ws_base(base: str, route: str) -> str:
|
||
if not route:
|
||
return base
|
||
return f"{base.rstrip('/')}{route}"
|
||
|
||
|
||
def _ensure_scheme(base: str, allowed: Tuple[str, ...]) -> None:
|
||
parsed = urlparse(base)
|
||
if not parsed.scheme or parsed.scheme.lower() not in allowed:
|
||
allowed_str = ", ".join(allowed)
|
||
raise ValueError(
|
||
f"🛑\u2009 Nilor-Nodes (comfyui_client): Base URL must start with one of [{allowed_str}]; got: {base!r}"
|
||
)
|
||
|
||
|
||
def _validate_bases(http_base: str, ws_base: str) -> None:
|
||
_ensure_scheme(http_base, ("http", "https"))
|
||
_ensure_scheme(ws_base, ("ws", "wss"))
|
||
|
||
|
||
def _quote_client_id(client_id: str) -> str:
|
||
return quote(client_id, safe="")
|
||
|
||
|
||
def _ws_url(base_ws: str, client_id: str) -> str:
|
||
return f"{base_ws.rstrip('/')}/ws?clientId={_quote_client_id(client_id)}"
|
||
|
||
|
||
# Bind helper methods to class namespace (private) without exposing publicly
|
||
ComfyUILocalClient._http_request_json = _http_request_json # type: ignore[attr-defined]
|
||
ComfyUILocalClient._join_http = lambda self, route: _join_http_base(self._base_url, route) # type: ignore[attr-defined]
|
||
ComfyUILocalClient._join_ws = lambda self, route: _join_ws_base(self._ws_url, route) # type: ignore[attr-defined]
|