Skip to content

Commit bd39bbf

Browse files
authored
[Partner Nodes] fix: respect Retry-After header (Comfy-Org#14234)
Signed-off-by: bigcat88 <bigcat88@icloud.com>
1 parent 5955ddf commit bd39bbf

2 files changed

Lines changed: 31 additions & 0 deletions

File tree

comfy_api_nodes/util/_helpers.py

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,8 @@
44
import re
55
import time
66
from collections.abc import Callable
7+
from datetime import datetime, timezone
8+
from email.utils import parsedate_to_datetime
79
from io import BytesIO
810

911
from yarl import URL
@@ -91,6 +93,32 @@ async def sleep_with_interrupt(
9193
await asyncio.sleep(min(1.0, end - now))
9294

9395

96+
def _retry_after_wait(value: str | None, fallback: float, max_wait: float) -> float:
97+
"""Delay before the next retry, honoring a server ``Retry-After`` header."""
98+
99+
seconds: float | None = None
100+
if value is not None:
101+
value = value.strip()
102+
if value.isascii() and value.isdigit():
103+
# delay-seconds form. The ASCII-digit guard keeps exotic Unicode "digit" characters away from float()
104+
# an all-digit string always converts (huge values become inf, never raising).
105+
seconds = float(value)
106+
elif value:
107+
# HTTP-date form. parsedate_to_datetime raises OverflowError (not a ValueError) on absurd years/offsets
108+
try:
109+
parsed = parsedate_to_datetime(value)
110+
except (TypeError, ValueError, OverflowError):
111+
parsed = None
112+
if parsed is not None:
113+
if parsed.tzinfo is None: # naive datetime: HTTP-date is UTC
114+
parsed = parsed.replace(tzinfo=timezone.utc)
115+
delta = (parsed - datetime.now(timezone.utc)).total_seconds()
116+
seconds = delta if delta > 0 else 0.0
117+
if seconds is None:
118+
return fallback
119+
return min(seconds, max_wait)
120+
121+
94122
def mimetype_to_extension(mime_type: str) -> str:
95123
"""Converts a MIME type to a file extension."""
96124
return mime_type.split("/")[-1].lower()

comfy_api_nodes/util/client.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121

2222
from . import request_logger
2323
from ._helpers import (
24+
_retry_after_wait,
2425
default_base_url,
2526
get_comfy_api_headers,
2627
get_node_id,
@@ -82,6 +83,7 @@ class _PollUIState:
8283

8384

8485
_RETRY_STATUS = {408, 500, 502, 503, 504} # status 429 is handled separately
86+
_MAX_RETRY_AFTER_WAIT = 150.0 # Cap a server Retry-After at this many seconds so a large hint can't block execution
8587
COMPLETED_STATUSES = ["succeeded", "succeed", "success", "completed", "finished", "done", "complete"]
8688
FAILED_STATUSES = ["cancelled", "canceled", "canceling", "fail", "failed", "error"]
8789
QUEUED_STATUSES = ["created", "queued", "queueing", "submitted", "initializing", "wait", "in_queue"]
@@ -747,6 +749,7 @@ async def _monitor(stop_evt: asyncio.Event, start_ts: float):
747749
should_retry = True
748750

749751
if should_retry:
752+
wait_time = _retry_after_wait(resp.headers.get("Retry-After"), wait_time, _MAX_RETRY_AFTER_WAIT)
750753
logging.warning(
751754
"HTTP %s %s -> %s. Waiting %.2fs (%s).",
752755
method,

0 commit comments

Comments
 (0)