From bd39bbf0678ebd31c972fd365733a8c729f2cd74 Mon Sep 17 00:00:00 2001
From: Alexander Piskun <13381981+bigcat88@users.noreply.github.com>
Date: Fri, 19 Jun 2026 11:32:56 +0300
Subject: [PATCH] [Partner Nodes] fix: respect Retry-After header (#14234)

Signed-off-by: bigcat88 <bigcat88@icloud.com>
---
 comfy_api_nodes/util/_helpers.py | 28 ++++++++++++++++++++++++++++
 comfy_api_nodes/util/client.py   |  3 +++
 2 files changed, 31 insertions(+)

diff --git a/comfy_api_nodes/util/_helpers.py b/comfy_api_nodes/util/_helpers.py
index 83cf7b001..6b8121cab 100644
--- a/comfy_api_nodes/util/_helpers.py
+++ b/comfy_api_nodes/util/_helpers.py
@@ -4,6 +4,8 @@ import os
 import re
 import time
 from collections.abc import Callable
+from datetime import datetime, timezone
+from email.utils import parsedate_to_datetime
 from io import BytesIO
 
 from yarl import URL
@@ -91,6 +93,32 @@ async def sleep_with_interrupt(
         await asyncio.sleep(min(1.0, end - now))
 
 
+def _retry_after_wait(value: str | None, fallback: float, max_wait: float) -> float:
+    """Delay before the next retry, honoring a server ``Retry-After`` header."""
+
+    seconds: float | None = None
+    if value is not None:
+        value = value.strip()
+        if value.isascii() and value.isdigit():
+            # delay-seconds form. The ASCII-digit guard keeps exotic Unicode "digit" characters away from float()
+            # an all-digit string always converts (huge values become inf, never raising).
+            seconds = float(value)
+        elif value:
+            # HTTP-date form. parsedate_to_datetime raises OverflowError (not a ValueError) on absurd years/offsets
+            try:
+                parsed = parsedate_to_datetime(value)
+            except (TypeError, ValueError, OverflowError):
+                parsed = None
+            if parsed is not None:
+                if parsed.tzinfo is None:  # naive datetime: HTTP-date is UTC
+                    parsed = parsed.replace(tzinfo=timezone.utc)
+                delta = (parsed - datetime.now(timezone.utc)).total_seconds()
+                seconds = delta if delta > 0 else 0.0
+    if seconds is None:
+        return fallback
+    return min(seconds, max_wait)
+
+
 def mimetype_to_extension(mime_type: str) -> str:
     """Converts a MIME type to a file extension."""
     return mime_type.split("/")[-1].lower()
diff --git a/comfy_api_nodes/util/client.py b/comfy_api_nodes/util/client.py
index adcde7bcb..66aab17f8 100644
--- a/comfy_api_nodes/util/client.py
+++ b/comfy_api_nodes/util/client.py
@@ -21,6 +21,7 @@ from server import PromptServer
 
 from . import request_logger
 from ._helpers import (
+    _retry_after_wait,
     default_base_url,
     get_comfy_api_headers,
     get_node_id,
@@ -82,6 +83,7 @@ class _PollUIState:
 
 
 _RETRY_STATUS = {408, 500, 502, 503, 504}  # status 429 is handled separately
+_MAX_RETRY_AFTER_WAIT = 150.0  # Cap a server Retry-After at this many seconds so a large hint can't block execution
 COMPLETED_STATUSES = ["succeeded", "succeed", "success", "completed", "finished", "done", "complete"]
 FAILED_STATUSES = ["cancelled", "canceled", "canceling", "fail", "failed", "error"]
 QUEUED_STATUSES = ["created", "queued", "queueing", "submitted", "initializing", "wait", "in_queue"]
@@ -747,6 +749,7 @@ async def _request_base(cfg: _RequestConfig, expect_binary: bool):
                         should_retry = True
 
                     if should_retry:
+                        wait_time = _retry_after_wait(resp.headers.get("Retry-After"), wait_time, _MAX_RETRY_AFTER_WAIT)
                         logging.warning(
                             "HTTP %s %s -> %s. Waiting %.2fs (%s).",
                             method,