nodes/silveroxides-ComfyUI_Gemini_Expanded_API
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Abort HTTP transport on timeout or interrupt and disable hidden SDK retries

Browse Source
This commit is contained in:
silveroxides
2026-10-24 15:16:02 +02:00
parent 357685d80b
commit 957672677f
3 changed files with 102 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
gemini_nodes.py
+32
View File
@@ -978,6 +978,21 @@ class SSL_GeminiTextPrompt(IO.ComfyNode):
def _check_interrupt(cls):
comfy.model_management.throw_exception_if_processing_interrupted()
@classmethod
def _abort_client_transport(cls, client: Any, client_key: Any = None) -> None:
if client_key is not None:
cls._client_cache.pop(client_key, None)
if client is None:
return
api_client = getattr(client, "_api_client", None)
if api_client is not None:
http_client = getattr(api_client, "_httpx_client", None)
if http_client is not None and hasattr(http_client, "close"):
try:
http_client.close()
except Exception:
pass
@classmethod
def _wait_for_result(cls, result_queue: "queue.Queue[Tuple[str, Any]]", timeout: float) -> Tuple[str, Any]:
deadline = time.monotonic() + float(timeout)
@@ -1366,8 +1381,11 @@ class SSL_GeminiTextPrompt(IO.ComfyNode):
# API call in background thread
start_time = time.time()
result_queue: "queue.Queue[Tuple[str, Any]]" = queue.Queue()
cancel_event = threading.Event()
def generate_once():
if cancel_event.is_set():
raise RuntimeError("Request cancelled")
# Apply the wire limit even on legacy unseeded image branches.
if generate_content_config.seed is not None and not -(2**31) <= generate_content_config.seed < 2**31:
generate_content_config.seed = cls._handle_seed(True, generate_content_config.seed)
@@ -1417,6 +1435,8 @@ class SSL_GeminiTextPrompt(IO.ComfyNode):
api_response = None
max_retries = 1
for attempt in range(max_retries):
if cancel_event.is_set():
return
try:
if use_seed and actual_seed is not None:
try:
@@ -1424,6 +1444,8 @@ class SSL_GeminiTextPrompt(IO.ComfyNode):
except Exception:
pass
response = generate_once()
if cancel_event.is_set():
return
if not (response.candidates and getattr(response.candidates[0].content, 'parts', None)):
finish_reason = "UNKNOWN"
if response.candidates:
@@ -1435,6 +1457,8 @@ class SSL_GeminiTextPrompt(IO.ComfyNode):
break
except Exception as e:
last_api_exception = e
if cancel_event.is_set():
return
if api_response is None:
result_queue.put(("error", last_api_exception))
return
@@ -1563,8 +1587,16 @@ class SSL_GeminiTextPrompt(IO.ComfyNode):
print(f"[WARNING] Invalid retry regex pattern: {regex_err}")
except queue.Empty:
cancel_event.set()
cls._abort_client_transport(client, client_key)
text_output = timeout_fallback_text or f"Gemini API request/processing timed out, waited {timeout} seconds."
except comfy.model_management.InterruptProcessingException:
cancel_event.set()
cls._abort_client_transport(client, client_key)
raise
except comfy.model_management.InterruptProcessingException:
raise
except Exception as e:
is_success = False
print(f"[ERROR] Unhandled error in generate method: {str(e)}")
interactions_nodes.py
+54
View File
@@ -803,6 +803,38 @@ class SSL_GeminiInteractionsTextPrompt(IO.ComfyNode):
def _check_interrupt(cls):
comfy.model_management.throw_exception_if_processing_interrupted()
@classmethod
def _configure_interactions_client(cls, client: Any, timeout: float) -> None:
try:
interactions = getattr(client, "interactions", None)
sdk_cfg = getattr(interactions, "sdk_configuration", None)
if sdk_cfg is not None:
sdk_cfg.timeout_ms = max(1, int(float(timeout) * 1000))
retry_cfg = getattr(sdk_cfg, "retry_config", None)
if retry_cfg is not None:
retry_cfg.strategy = "none"
retry_cfg.max_retries = 0
retry_cfg.retry_connection_errors = False
except Exception:
pass
@classmethod
def _abort_client_transport(cls, client: Any, client_key: Any = None) -> None:
if client_key is not None:
cls._client_cache.pop(client_key, None)
if client is None:
return
for owner in (getattr(client, "_api_client", None), getattr(getattr(client, "_interactions", None), "sdk_configuration", None)):
if owner is None:
continue
for attr in ("_httpx_client", "client"):
http_client = getattr(owner, attr, None)
if http_client is not None and hasattr(http_client, "close"):
try:
http_client.close()
except Exception:
pass
@classmethod
def _wait_for_result(cls, result_queue: "queue.Queue[Tuple[str, Any]]", timeout: float) -> Tuple[str, Any]:
deadline = time.monotonic() + float(timeout)
@@ -1203,10 +1235,14 @@ class SSL_GeminiInteractionsTextPrompt(IO.ComfyNode):
generation_config["thinking_summaries"] = "auto"
# Background worker queue for API call
cls._configure_interactions_client(client, timeout)
start_time = time.time()
result_queue: "queue.Queue[Tuple[str, Any]]" = queue.Queue()
cancel_event = threading.Event()
def create_interaction_once(current_seed, target_previous_id):
if cancel_event.is_set():
raise RuntimeError("Request cancelled")
call_gen_config = dict(generation_config)
if use_seed and current_seed is not None:
call_gen_config["seed"] = int(current_seed)
@@ -1235,12 +1271,18 @@ class SSL_GeminiInteractionsTextPrompt(IO.ComfyNode):
api_response = None
max_attempts = 1
for attempt in range(max_attempts):
if cancel_event.is_set():
return
try:
current_seed = (actual_seed + attempt) if (use_seed and actual_seed is not None) else None
response = create_interaction_once(current_seed, effective_previous_id)
if cancel_event.is_set():
return
api_response = response
break
except Exception as e:
if cancel_event.is_set():
return
err_str = str(e).lower()
# If previous interaction was specified/auto-used but expired or not found, fall back to new interaction
if effective_previous_id and ("not found" in err_str or "expired" in err_str or "404" in err_str or "invalid" in err_str):
@@ -1248,12 +1290,16 @@ class SSL_GeminiInteractionsTextPrompt(IO.ComfyNode):
effective_previous_id = ""
try:
response = create_interaction_once(current_seed, "")
if cancel_event.is_set():
return
api_response = response
break
except Exception as retry_err:
last_api_exception = retry_err
else:
last_api_exception = e
if cancel_event.is_set():
return
if api_response is None:
result_queue.put(("error", last_api_exception))
return
@@ -1424,8 +1470,16 @@ class SSL_GeminiInteractionsTextPrompt(IO.ComfyNode):
print(f"[WARNING] Invalid retry regex pattern: {regex_err}")
except queue.Empty:
cancel_event.set()
cls._abort_client_transport(client, client_key)
text_output = timeout_fallback_text or f"Gemini API request/processing timed out, waited {timeout} seconds."
except comfy.model_management.InterruptProcessingException:
cancel_event.set()
cls._abort_client_transport(client, client_key)
raise
except comfy.model_management.InterruptProcessingException:
raise
except Exception as e:
is_success = False
print(f"[ERROR] Unhandled error in generate method: {str(e)}")
tests/test_interactions_nodes.py
+16
View File
@@ -471,24 +471,40 @@ def test_retry_pattern_triggers_retry(mock_interactions_client):
def test_timeout_fallback(mock_interactions_client, monkeypatch):
closed = {"count": 0}
retry_cfg = SimpleNamespace(strategy="attempt-count-backoff", max_retries=3, retry_connection_errors=True)
class FakeHttpClient:
def close(self):
closed["count"] += 1
def slow_create(**kwargs):
import time
time.sleep(1.0)
return SimpleNamespace(id="slow", output_text="slow")
class SlowInteractions:
def __init__(self):
self.sdk_configuration = SimpleNamespace(timeout_ms=None, retry_config=retry_cfg, client=FakeHttpClient())
def create(self, **kwargs):
return slow_create(**kwargs)
class SlowClient:
def __init__(self, **kwargs):
self._api_client = SimpleNamespace(_httpx_client=FakeHttpClient())
self.interactions = SlowInteractions()
self._interactions = self.interactions
monkeypatch.setattr(interactions_nodes.genai, "Client", SlowClient)
args = _execute_kwargs(timeout=0.1, timeout_fallback_text="Custom timeout message")
output = interactions_nodes.SSL_GeminiInteractionsTextPrompt.execute(**args)
assert output[0] == "Custom timeout message"
assert retry_cfg.strategy == "none"
assert retry_cfg.max_retries == 0
assert closed["count"] >= 1
assert len(interactions_nodes.SSL_GeminiInteractionsTextPrompt._client_cache) == 0
def test_interactions_api_key_config_with_cache():