Fix ctrl c not working during startup (#4155)

openhands/llm/retry_mixin.py

@@ -5,9 +5,8 @@ from tenacity import (
     wait_exponential,
 )
 
-from openhands.core.exceptions import OperationCancelled
 from openhands.core.logger import openhands_logger as logger
-from openhands.runtime.utils.shutdown_listener import should_exit
+from openhands.utils.tenacity_stop import stop_if_should_exit
 
 
 class RetryMixin:
@@ -32,7 +31,7 @@ class RetryMixin:
 
         return retry(
             before_sleep=self.log_retry_attempt,
-            stop=stop_after_attempt(num_retries),
+            stop=stop_after_attempt(num_retries) | stop_if_should_exit(),
             reraise=True,
             retry=(retry_if_exception_type(retry_exceptions)),
             wait=wait_exponential(
@@ -44,8 +43,6 @@ class RetryMixin:
 
     def log_retry_attempt(self, retry_state):
         """Log retry attempts."""
-        if should_exit():
-            raise OperationCancelled('Operation cancelled.')  # exits the @retry loop
         exception = retry_state.outcome.exception()
         logger.error(
             f'{exception}. Attempt #{retry_state.attempt_number} | You can customize retry values in the configuration.',

openhands/memory/memory.py

@@ -11,6 +11,7 @@ from tenacity import (
 from openhands.core.config import LLMConfig
 from openhands.core.logger import openhands_logger as logger
 from openhands.core.utils import json
+from openhands.utils.tenacity_stop import stop_if_should_exit
 
 try:
     import chromadb
@@ -50,7 +51,7 @@ if LLAMA_INDEX_AVAILABLE:
 
     @retry(
         reraise=True,
-        stop=stop_after_attempt(num_retries),
+        stop=stop_after_attempt(num_retries) | stop_if_should_exit(),
         wait=wait_random_exponential(min=retry_min_wait, max=retry_max_wait),
         retry=retry_if_exception_type(
             (RateLimitError, APIConnectionError, InternalServerError)

openhands/runtime/browser/browser_env.py

@@ -16,6 +16,7 @@ from PIL import Image
 
 from openhands.core.exceptions import BrowserInitException
 from openhands.core.logger import openhands_logger as logger
+from openhands.utils.tenacity_stop import stop_if_should_exit
 from openhands.runtime.utils.shutdown_listener import should_continue, should_exit
 
 BROWSER_EVAL_GET_GOAL_ACTION = 'GET_EVAL_GOAL'
@@ -52,7 +53,7 @@ class BrowserEnv:
 
     @tenacity.retry(
         wait=tenacity.wait_fixed(1),
-        stop=tenacity.stop_after_attempt(5),
+        stop=tenacity.stop_after_attempt(5) | stop_if_should_exit(),
         retry=tenacity.retry_if_exception_type(BrowserInitException),
     )
     def init_browser(self):

openhands/runtime/client/runtime.py

@@ -36,6 +36,7 @@ from openhands.runtime.plugins import PluginRequirement
 from openhands.runtime.runtime import Runtime
 from openhands.runtime.utils import find_available_tcp_port
 from openhands.runtime.utils.runtime_build import build_runtime_image
+from openhands.utils.tenacity_stop import stop_if_should_exit
 
 
 class LogBuffer:
@@ -201,7 +202,7 @@ class EventStreamRuntime(Runtime):
             raise ex
 
     @tenacity.retry(
-        stop=tenacity.stop_after_attempt(5),
+        stop=tenacity.stop_after_attempt(5) | stop_if_should_exit(),
         wait=tenacity.wait_exponential(multiplier=1, min=4, max=60),
     )
     def _init_container(
@@ -322,7 +323,7 @@ class EventStreamRuntime(Runtime):
             )
 
     @tenacity.retry(
-        stop=tenacity.stop_after_attempt(10),
+        stop=tenacity.stop_after_delay(120) | stop_if_should_exit(),
         wait=tenacity.wait_exponential(multiplier=2, min=1, max=20),
         reraise=(ConnectionRefusedError,),
     )

openhands/runtime/remote/runtime.py

@@ -36,6 +36,7 @@ from openhands.events.serialization.action import ACTION_TYPE_TO_CLASS
 from openhands.runtime.builder.remote import RemoteRuntimeBuilder
 from openhands.runtime.plugins import PluginRequirement
 from openhands.runtime.runtime import Runtime
+from openhands.utils.tenacity_stop import stop_if_should_exit
 from openhands.runtime.utils.request import (
     DEFAULT_RETRY_EXCEPTIONS,
     is_404_error,
@@ -191,7 +192,7 @@ class RemoteRuntime(Runtime):
         ), 'Runtime URL is not set. This should never happen.'
 
     @retry(
-        stop=stop_after_attempt(10),
+        stop=stop_after_attempt(10) | stop_if_should_exit(),
         wait=wait_exponential(multiplier=1, min=4, max=60),
         retry=retry_if_exception_type(RuntimeError),
         reraise=True,

openhands/runtime/utils/request.py

@@ -10,6 +10,8 @@ from tenacity import (
     wait_exponential,
 )
 
+from openhands.utils.tenacity_stop import stop_if_should_exit
+
 
 def is_server_error(exception):
     return (
@@ -50,7 +52,7 @@ def send_request(
     kwargs["timeout"] = timeout
 
     @retry(
-        stop=stop_after_delay(timeout),
+        stop=stop_after_delay(timeout) | stop_if_should_exit(),
         wait=wait_exponential(multiplier=1, min=4, max=60),
         retry=retry_condition,
         reraise=True,

openhands/utils/tenacity_stop.py

@@ -0,0 +1,11 @@
+from tenacity import RetryCallState
+from tenacity.stop import stop_base
+
+from openhands.runtime.utils.shutdown_listener import should_exit
+
+
+class stop_if_should_exit(stop_base):
+    """Stop if the should_exit flag is set."""
+
+    def __call__(self, retry_state: 'RetryCallState') -> bool:
+        return should_exit()