Add OpenRouter request timeout

README.md

@@ -95,8 +95,9 @@ To keep your API key secure, use one of the following methods:
 - **web_search**: Enable web search capability by appending `:online` to the model ID. This costs $4 per 1000 queries and automatically uses your openrouter balance.
 - **cheapest**: Route to the cheapest provider by appending `:floor` to the model ID (enabled by default).
 - **fastest**: Route to the fastest provider by appending `:nitro` to the model ID (disabled by default).
-- **temperature**: Controls the randomness of the model's output (0.0 to 2.0).
-- **chat_mode**: Enable conversation mode to maintain context across messages (disabled by default).
+- **temperature**: Controls the randomness of the model's output (0.0 to 2.0).
+- **chat_mode**: Enable conversation mode to maintain context across messages (disabled by default).
+- **request_timeout**: Maximum time, in seconds, to wait for the main OpenRouter completion request before returning an error (default: 120).
 
 #### Optional Inputs:
 

node.py

@@ -26,9 +26,12 @@ class OpenRouterNode:
       3) "Credits": a string showing your remaining OpenRouter account balance
     """
 
-    models_cache = None
-    last_fetch_time = 0
-    cache_duration = 3600  # Cache duration in seconds (1 hour)
+    models_cache = None
+    last_fetch_time = 0
+    cache_duration = 3600  # Cache duration in seconds (1 hour)
+    default_request_timeout = 120
+    min_request_timeout = 1
+    max_request_timeout = 3600
 
     def __init__(self):
         self.chat_manager = ChatSessionManager()
@@ -116,11 +119,18 @@ class OpenRouterNode:
                     "round": 0.01,
                 }),
                  "pdf_engine": (["auto", "mistral-ocr", "pdf-text"], {"default": "auto"}),
-                "chat_mode": ("BOOLEAN", {"default": False}),
-            },
-            "optional": {
-                "pdf_data": (PDF_DATA_TYPE,), # Use '*' and check structure in generate_response
-                "user_message_input": ("STRING", {"forceInput": True}),
+                "chat_mode": ("BOOLEAN", {"default": False}),
+                "request_timeout": ("INT", {
+                    "default": cls.default_request_timeout,
+                    "min": cls.min_request_timeout,
+                    "max": cls.max_request_timeout,
+                    "step": 1,
+                    "display": "number",
+                }),
+            },
+            "optional": {
+                "pdf_data": (PDF_DATA_TYPE,), # Use '*' and check structure in generate_response
+                "user_message_input": ("STRING", {"forceInput": True}),
             }
         }
 
@@ -139,7 +149,7 @@ class OpenRouterNode:
         if (cls.models_cache is None) or (current_time - cls.last_fetch_time > cls.cache_duration):
             url = "https://openrouter.ai/api/v1/models"
             try:
-                response = requests.get(url)
+                response = requests.get(url, timeout=cls.default_request_timeout)
                 response.raise_for_status()
                 models = response.json()["data"]
                 # Filter for models that support chat completions if needed, but API handles this
@@ -153,17 +163,27 @@ class OpenRouterNode:
                     cls.models_cache = ["error_fetching_models", "google/gemma-3-27b-it", "openai/gpt-4o"] # Example fallbacks
         return cls.models_cache if cls.models_cache else ["error_fetching_models"] # Ensure it's never empty
 
-    def validate_temperature(self, temperature):
-        """
-        Validates and converts temperature value to float within acceptable range.
-        """
-        try:
+    def validate_temperature(self, temperature):
+        """
+        Validates and converts temperature value to float within acceptable range.
+        """
+        try:
             temp = float(temperature)
             return max(0.0, min(2.0, temp))  # Clamp between 0.0 and 2.0
-        except (ValueError, TypeError):
-            return 1.0  # Return default if conversion fails
-
-    def fetch_credits(self, api_key):
+        except (ValueError, TypeError):
+            return 1.0  # Return default if conversion fails
+
+    def validate_request_timeout(self, request_timeout):
+        """
+        Validates and converts request timeout to seconds within an acceptable range.
+        """
+        try:
+            timeout = int(request_timeout)
+            return max(self.min_request_timeout, min(self.max_request_timeout, timeout))
+        except (ValueError, TypeError):
+            return self.default_request_timeout
+
+    def fetch_credits(self, api_key, timeout=None):
         """
         Fetches the user's credits information from the OpenRouter API.
         Returns a formatted string with remaining credits.
@@ -180,8 +200,9 @@ class OpenRouterNode:
             "X-Title": "ComfyUI OpenRouter LLM Node",
         }
 
-        try:
-            response = requests.get(url, headers=headers)
+        try:
+            validated_timeout = self.validate_request_timeout(timeout)
+            response = requests.get(url, headers=headers, timeout=validated_timeout)
             response.raise_for_status()
 
             result = response.json()
@@ -205,10 +226,10 @@ class OpenRouterNode:
         except json.JSONDecodeError:
              return "Error fetching credits: Could not decode JSON response."
 
-    def generate_response(self, api_key, system_prompt, user_message_box, model,
-                         web_search, cheapest, fastest, temperature, pdf_engine, chat_mode,
-                         aspect_ratio="auto", image_resolution="1K", seed=0,
-                         pdf_data=None, user_message_input=None, reasoning_effort="none", **kwargs):
+    def generate_response(self, api_key, system_prompt, user_message_box, model,
+                         web_search, cheapest, fastest, temperature, pdf_engine, chat_mode,
+                         request_timeout=120, aspect_ratio="auto", image_resolution="1K", seed=0,
+                         pdf_data=None, user_message_input=None, reasoning_effort="none", **kwargs):
         """
         Sends a completion request to the OpenRouter chat completion endpoint.
         Handles text, optional image, and optional PDF inputs.
@@ -236,8 +257,9 @@ class OpenRouterNode:
             "X-Title": "ComfyUI OpenRouter LLM Node",
         }
 
-        # Validate and convert temperature
-        validated_temp = self.validate_temperature(temperature)
+        # Validate and convert temperature
+        validated_temp = self.validate_temperature(temperature)
+        validated_timeout = self.validate_request_timeout(request_timeout)
 
         # Decide whether to use user_message_input or user_message_box
         user_text = user_message_input if user_message_input is not None and user_message_input.strip() else user_message_box
@@ -389,10 +411,10 @@ class OpenRouterNode:
 
         # --- Make API Call and Process Response ---
         try:
-            start_time = time.time()
-            response = requests.post(url, headers=headers, json=data)
-            response.raise_for_status() # Raises HTTPError for bad responses (4xx or 5xx)
-            end_time = time.time()
+            start_time = time.time()
+            response = requests.post(url, headers=headers, json=data, timeout=validated_timeout)
+            response.raise_for_status() # Raises HTTPError for bad responses (4xx or 5xx)
+            end_time = time.time()
 
             result = response.json()
             # Debug: print truncated response to see what OpenRouter returned
@@ -487,7 +509,7 @@ class OpenRouterNode:
 
 
             # Fetch credits information AFTER the main request
-            credits_text = self.fetch_credits(api_key)
+            credits_text = self.fetch_credits(api_key, timeout=validated_timeout)
 
             # Save conversation in chat mode
             if chat_mode and session_path:
@@ -625,11 +647,11 @@ class OpenRouterNode:
             return max(1, round(len(text) / 4))
 
 
-    @classmethod
-    def IS_CHANGED(cls, api_key, system_prompt, user_message_box, model,
-                   web_search, cheapest, fastest, temperature, pdf_engine, chat_mode,
-                   aspect_ratio="auto", image_resolution="1K", seed=0,
-                   pdf_data=None, user_message_input=None, **kwargs):
+    @classmethod
+    def IS_CHANGED(cls, api_key, system_prompt, user_message_box, model,
+                   web_search, cheapest, fastest, temperature, pdf_engine, chat_mode,
+                   request_timeout=120, aspect_ratio="auto", image_resolution="1K", seed=0,
+                   pdf_data=None, user_message_input=None, **kwargs):
         """
         Check if any input that affects the output has changed.
         Includes hashing image and PDF data.
@@ -675,8 +697,14 @@ class OpenRouterNode:
         try:
             temp_float = float(temperature) if isinstance(temperature, (str, int, float)) else 1.0
             temp_float = max(0.0, min(2.0, temp_float))
-        except (ValueError, TypeError):
-            temp_float = 1.0
+        except (ValueError, TypeError):
+            temp_float = 1.0
+
+        try:
+            timeout_int = int(request_timeout)
+            timeout_int = max(cls.min_request_timeout, min(cls.max_request_timeout, timeout_int))
+        except (ValueError, TypeError):
+            timeout_int = cls.default_request_timeout
 
 
         # Combine all relevant inputs into a tuple for comparison
@@ -684,9 +712,9 @@ class OpenRouterNode:
         # Note: api_key here is the UI value only. Keys resolved from the LLM_KEY
         # env var or openrouter_api_key.json are intentionally NOT part of the
         # cache key — they're treated as user environment, not workflow inputs.
-        return (api_key, system_prompt, user_message_box, model,
-                web_search, cheapest, fastest, temp_float, pdf_engine, chat_mode,
-                aspect_ratio, image_resolution, seed, tuple(image_hashes), pdf_hash, user_message_input)
+        return (api_key, system_prompt, user_message_box, model,
+                web_search, cheapest, fastest, temp_float, pdf_engine, chat_mode,
+                timeout_int, aspect_ratio, image_resolution, seed, tuple(image_hashes), pdf_hash, user_message_input)
 
 # Node class mappings
 NODE_CLASS_MAPPINGS = {

tests/__init__.py

@@ -0,0 +1 @@
+

tests/test_request_timeout.py

@@ -0,0 +1,138 @@
+import importlib.util
+import sys
+import types
+import unittest
+from pathlib import Path
+from unittest.mock import Mock, patch
+
+
+class FakeRequestException(Exception):
+    def __init__(self, message="", response=None):
+        super().__init__(message)
+        self.response = response
+
+
+class FakeTimeout(FakeRequestException):
+    pass
+
+
+def install_dependency_stubs():
+    requests_module = types.ModuleType("requests")
+    exceptions = types.SimpleNamespace(
+        RequestException=FakeRequestException,
+        Timeout=FakeTimeout,
+    )
+    requests_module.exceptions = exceptions
+    requests_module.get = Mock()
+    requests_module.post = Mock()
+    sys.modules["requests"] = requests_module
+
+    torch_module = types.ModuleType("torch")
+    torch_module.float32 = "float32"
+    torch_module.Tensor = type("Tensor", (), {})
+    torch_module.zeros = Mock(return_value="placeholder-image")
+    sys.modules["torch"] = torch_module
+
+    tiktoken_module = types.ModuleType("tiktoken")
+    tiktoken_module.get_encoding = Mock()
+    sys.modules["tiktoken"] = tiktoken_module
+
+    pil_module = types.ModuleType("PIL")
+    image_module = types.ModuleType("PIL.Image")
+    image_module.open = Mock()
+    image_module.fromarray = Mock()
+    pil_module.Image = image_module
+    sys.modules["PIL"] = pil_module
+    sys.modules["PIL.Image"] = image_module
+
+    sys.modules.setdefault("numpy", types.ModuleType("numpy"))
+
+
+def load_node_module():
+    install_dependency_stubs()
+    root = Path(__file__).resolve().parents[1]
+    package_name = "comfy_openrouter_node_test"
+
+    package = types.ModuleType(package_name)
+    package.__path__ = [str(root)]
+    sys.modules[package_name] = package
+
+    spec = importlib.util.spec_from_file_location(
+        f"{package_name}.node",
+        root / "node.py",
+        submodule_search_locations=[str(root)],
+    )
+    module = importlib.util.module_from_spec(spec)
+    sys.modules[spec.name] = module
+    spec.loader.exec_module(module)
+    return module
+
+
+class RequestTimeoutTests(unittest.TestCase):
+    def setUp(self):
+        self.node_module = load_node_module()
+        self.node = self.node_module.OpenRouterNode()
+        self.node.fetch_credits = Mock(return_value="Remaining: $1.000")
+        self.node.count_tokens = Mock(return_value=1)
+
+    def call_generate_response(self, request_timeout=120):
+        return self.node.generate_response(
+            api_key="test-key",
+            system_prompt="system",
+            user_message_box="hello",
+            model="openai/gpt-4o",
+            web_search=False,
+            cheapest=False,
+            fastest=False,
+            temperature=1.0,
+            pdf_engine="auto",
+            chat_mode=False,
+            request_timeout=request_timeout,
+        )
+
+    def test_main_openrouter_request_uses_configured_timeout(self):
+        response = Mock()
+        response.raise_for_status = Mock()
+        response.json.return_value = {
+            "choices": [{"message": {"content": "done"}}],
+            "usage": {"prompt_tokens": 2, "completion_tokens": 3},
+        }
+        self.node_module.requests.post.return_value = response
+
+        result = self.call_generate_response(request_timeout=45)
+
+        self.assertEqual(result[0], "done")
+        self.node_module.requests.post.assert_called_once()
+        self.assertEqual(self.node_module.requests.post.call_args.kwargs["timeout"], 45)
+
+    def test_timeout_exception_returns_clear_error(self):
+        self.node_module.requests.post.side_effect = FakeTimeout("request timed out")
+
+        result = self.call_generate_response(request_timeout=2)
+
+        self.assertIn("API Request Error", result[0])
+        self.assertIn("request timed out", result[0])
+        self.assertEqual(result[2], "Stats N/A due to error")
+        self.assertEqual(result[3], "Credits N/A due to error")
+
+    def test_fetch_credits_uses_configured_timeout(self):
+        self.node.fetch_credits = self.node_module.OpenRouterNode().fetch_credits
+        response = Mock()
+        response.raise_for_status = Mock()
+        response.json.return_value = {
+            "data": {
+                "total_credits": 10.0,
+                "total_usage": 3.25,
+            }
+        }
+        self.node_module.requests.get.return_value = response
+
+        credits = self.node.fetch_credits("test-key", timeout=45)
+
+        self.assertEqual(credits, "Remaining: $6.750")
+        self.node_module.requests.get.assert_called_once()
+        self.assertEqual(self.node_module.requests.get.call_args.kwargs["timeout"], 45)
+
+
+if __name__ == "__main__":
+    unittest.main()