Comfy-Org-comfy-remote-nodes/protocol.py

"""Vendored copy of the parts of ``comfy_rnp_protocol`` the client needs.

The full protocol package lives in ``comfy-rnp-server/comfy_rnp_protocol``
and is the source of truth for these constants and helpers — this file
is a thin mirror so the ComfyUI extension can ship without depending on
the server package being on ``PYTHONPATH``. Keep it byte-compatible with
the upstream module.
"""
from __future__ import annotations

import base64
from typing import Any

PROTOCOL_NAME = "comfy-rnp"
PROTOCOL_VERSION = "1.0"
SCHEMA_HASH_ALGORITHM = "sha256"
DEFAULT_MAX_INLINE_PAYLOAD_BYTES = 8 * 1024 * 1024
DEFAULT_POLL_INTERVAL_S = 1.5

# Defaults for descriptor.remote.execution fields the client must honour.
# Values mirror comfy_api_nodes.util.client.poll_op_raw so RNP-routed
# nodes behave identically to a direct upstream call when the descriptor
# leaves a field unset.
#
# Lifetime budget design: the wire format declares ``soft_timeout_s``
# (server-side stuck threshold) and ``hard_timeout_s`` (server-side
# hard kill) as durations. The client's poll-loop cap is a derived
# value — ``max_poll_attempts = ceil((hard_timeout_s + DEFAULT_POLL_GRACE_S)
# / poll_interval_s)`` — not a separate wire field, so changing the
# poll cadence can never silently shrink or stretch the contract.
# Industry precedent: OAuth2 device flow §3.5, Kubernetes
# ``--timeout``, gRPC deadlines all express bounds as durations.
DEFAULT_SOFT_TIMEOUT_S = 240.0
DEFAULT_HARD_TIMEOUT_S = 300.0
DEFAULT_POLL_GRACE_S = 30.0
DEFAULT_TIMEOUT_PER_POLL_S = 120.0
DEFAULT_MAX_RETRIES_PER_POLL = 10
DEFAULT_RETRY_DELAY_PER_POLL_S = 1.0
DEFAULT_RETRY_BACKOFF_PER_POLL = 1.4
DEFAULT_CANCEL_TIMEOUT_S = 10.0
DEFAULT_SUBMIT_MAX_RETRIES = 2
DEFAULT_SUBMIT_RETRY_BACKOFF = 2.0
DEFAULT_SUBMIT_RETRY_DELAY_S = 1.0


class IdempotencyMode:
    CLIENT_KEY = "client_key"
    NONE = "none"


class ExecutionMode:
    REQUEST_RESPONSE = "request_response"
    ASYNC_POLLING = "async_polling"


class TaskStatus:
    PENDING = "pending"
    RUNNING = "running"
    DONE = "done"
    ERROR = "error"
    CANCELLED = "cancelled"


TERMINAL_TASK_STATES = frozenset({
    TaskStatus.DONE,
    TaskStatus.ERROR,
    TaskStatus.CANCELLED,
})


class ErrorCode:
    INPUT_INVALID = "INPUT_INVALID"
    AUTH_FAILED = "AUTH_FAILED"
    INSUFFICIENT_CREDITS = "INSUFFICIENT_CREDITS"
    RATE_LIMITED = "RATE_LIMITED"
    PROVIDER_UNAVAILABLE = "PROVIDER_UNAVAILABLE"
    PAYLOAD_TOO_LARGE = "PAYLOAD_TOO_LARGE"
    VERSION_INCOMPATIBLE = "VERSION_INCOMPATIBLE"
    SCHEMA_HASH_MISMATCH = "SCHEMA_HASH_MISMATCH"
    TIMEOUT = "TIMEOUT"
    INTERNAL = "INTERNAL"
    NOT_FOUND = "NOT_FOUND"
    NOT_IMPLEMENTED = "NOT_IMPLEMENTED"
    # 404 from a poll/cancel that carried X-RNP-Idempotency-Key — the
    # server lost the task (typically a restart). The client replays
    # the original submit with the same key to spawn a fresh task.
    TASK_LOST = "TASK_LOST"
    # Server-side back-pressure / operator drain — both ride 503 with
    # Retry-After. SERVER_BUSY = transient saturation, MAINTENANCE =
    # operator action.
    SERVER_BUSY = "SERVER_BUSY"
    MAINTENANCE = "MAINTENANCE"


class Header:
    PROTOCOL_VERSION = "X-RNP-Protocol-Version"
    CLIENT_VERSION = "X-RNP-Client-Version"
    CLIENT_CAPABILITIES = "X-RNP-Client-Capabilities"
    SCHEMA_HASH = "X-RNP-Schema-Hash"
    IDEMPOTENCY_KEY = "X-RNP-Idempotency-Key"


class Capability:
    IMAGE_PNG_BASE64 = "image:png_base64"
    IMAGE_JPEG_BASE64 = "image:jpeg_base64"
    IMAGE_RAW_CHW_UINT8 = "image:raw_b64_chw_uint8"
    # Multi-frame IMAGE envelope: a single value carries N PNG-base64
    # frames sharing the same H/W/C, decoded into a [B,H,W,C] tensor on
    # this client. Mirrors the AUDIO sweep precedent (parallel encodings
    # rather than extending ``png_base64``) so legacy server emits still
    # decode correctly when the new token isn't advertised.
    IMAGE_PNG_BASE64_BATCH = "image:png_base64_batch"
    # Externalised counterpart of ``image:png_base64_batch``: same
    # multi-frame IMAGE envelope shape, but every frame's PNG bytes
    # are uploaded out-of-band via the presigned-PUT path (see
    # ``serialization.maybe_externalize``) and the envelope carries
    # a parallel ``uris`` list instead of inline ``frames``. Required
    # because the original ``maybe_externalize`` only inspected the
    # singleton ``data`` field, so a batch of N PNGs whose summed
    # base64 payload crossed the 8 MiB inline cap silently rode the
    # wire as a 30+ MiB POST body. Per-frame URIs (NOT a tar/zip
    # blob) keep this composable with the existing per-frame fetch
    # helpers and let the server fetch in parallel. Decoder dispatch
    # keys off the ``encoding`` string — ``frames`` vs ``uris`` is a
    # payload field, not a shape selector.
    IMAGE_PNG_BASE64_BATCH_URI = "image:png_base64_batch_uri"
    VIDEO_MP4_INLINE = "video:mp4_inline"
    VIDEO_MP4_BASE64 = "video:mp4_base64"
    AUDIO_MP3_INLINE = "audio:mp3_inline"
    AUDIO_MP3_BASE64 = "audio:mp3_base64"
    AUDIO_WAV_BASE64 = "audio:wav_base64"
    AUDIO_WAV_INLINE = "audio:wav_inline"
    MASK_PNG_BASE64 = "mask:png_base64"
    SVG_XML_BASE64 = "svg:svg_xml_base64"
    SCHEMA_V3 = "schema:v3"
    ASYNC_EXECUTE = "execute:async"
    # Pass-through custom-IO bucket: when an input/output type string
    # is not in the standard primitive/media set, the proxy_node parser
    # falls back to ``IO.Custom(<type>)`` so partner helper-config nodes
    # (OpenAIInputFiles, OpenAIChatConfig, GeminiInputFiles, Recraft
    # style/color/controls helpers, future partner helper-config types)
    # round-trip without per-type encoder plumbing. The descriptor's
    # original io_type string is preserved on the V3 input so node-to-
    # node connections only chain when the strings match. Encoder
    # dispatch passes opaque dict/list/string values through
    # ``_encode_one`` unchanged — they're already JSON-serialisable.
    IO_OPAQUE = "io:opaque"


HEAVY_TYPES = frozenset({"image", "video", "audio", "mask"})


def is_envelope(value: Any) -> bool:
    return (
        isinstance(value, dict)
        and isinstance(value.get("type"), str)
        and value.get("type") in HEAVY_TYPES
        and isinstance(value.get("encoding"), str)
    )


def decode_envelope_data(envelope: dict[str, Any]) -> bytes:
    """Decode the inline base64 ``data`` of an envelope.

    URI envelopes are resolved out-of-band by the caller (see
    ``proxy_node._deserialize_output``); by the time they reach this
    function the bytes have already been fetched and re-stamped onto
    ``data``.
    """
    if "data" not in envelope:
        raise ValueError(
            "envelope has no inline 'data' — caller must resolve "
            f"{sorted(k for k in envelope.keys() if k != 'data')!r} first"
        )
    return base64.b64decode(envelope["data"])


class RnpProtocolError(Exception):
    """Raised when a server response can't be parsed as RNP/1, or when
    a structured-error response is returned. Carries the parsed error
    envelope so the caller can surface ``user_facing`` messages.
    """

    def __init__(
        self,
        message: str,
        *,
        code: str = ErrorCode.INTERNAL,
        status: int | None = None,
        user_facing: bool = False,
        retryable: bool = False,
        details: dict[str, Any] | None = None,
    ) -> None:
        super().__init__(message)
        self.code = code
        self.status = status
        self.user_facing = user_facing
        self.retryable = retryable
        self.details = details

    @classmethod
    def from_response_body(
        cls, body: dict[str, Any] | None, status: int
    ) -> "RnpProtocolError":
        err = (body or {}).get("error") if isinstance(body, dict) else None
        if not isinstance(err, dict):
            return cls(
                f"Unstructured error response (HTTP {status})",
                code=ErrorCode.INTERNAL,
                status=status,
            )
        return cls(
            err.get("message") or f"HTTP {status}",
            code=err.get("code") or ErrorCode.INTERNAL,
            status=status,
            user_facing=bool(err.get("user_facing")),
            retryable=bool(err.get("retryable")),
            details=err.get("details") if isinstance(err.get("details"), dict) else None,
        )