diff --git a/gemini_nodes.py b/gemini_nodes.py index 087d2bc..14d985f 100644 --- a/gemini_nodes.py +++ b/gemini_nodes.py @@ -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)}") diff --git a/interactions_nodes.py b/interactions_nodes.py index 26f000a..9d9f1dc 100644 --- a/interactions_nodes.py +++ b/interactions_nodes.py @@ -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)}") diff --git a/tests/test_interactions_nodes.py b/tests/test_interactions_nodes.py index f63591c..7f47f8d 100644 --- a/tests/test_interactions_nodes.py +++ b/tests/test_interactions_nodes.py @@ -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():