Updated remove logic

openhands/runtime/builder/remote.py

@@ -18,11 +18,14 @@
 class RemoteRuntimeBuilder(RuntimeBuilder):
     """This class interacts with the remote Runtime API for building and managing container images."""
 
-    def __init__(self, api_url: str, api_key: str):
+    def __init__(
+        self, api_url: str, api_key: str, session: requests.Session | None = None
+    ):
         self.api_url = api_url
         self.api_key = api_key
-        self.session = requests.Session()
+        self.session = session or requests.Session()
         self.session.headers.update({'X-API-Key': self.api_key})
+        self._closed = False
 
     def build(
         self,
@@ -74,8 +77,8 @@ def build(
 
         # Poll /build_status until the build is complete
         start_time = time.time()
-        timeout = 30 * 60  # 20 minutes in seconds
-        while should_continue():
+        timeout = 30 * 60  # 30 minutes in seconds
+        while should_continue() and not self._closed:
             if time.time() - start_time > timeout:
                 logger.error('Build timed out after 30 minutes')
                 raise AgentRuntimeBuildError('Build timed out after 30 minutes')
@@ -116,9 +119,7 @@ def build(
             # Wait before polling again
             sleep_if_should_continue(30)
 
-        raise AgentRuntimeBuildError(
-            'Build interrupted (likely received SIGTERM or SIGINT).'
-        )
+        raise AgentRuntimeBuildError('Build interrupted')
 
     def image_exists(self, image_name: str, pull_from_repo: bool = True) -> bool:
         """Checks if an image exists in the remote registry using the /image_exists endpoint."""
@@ -148,3 +149,6 @@ def image_exists(self, image_name: str, pull_from_repo: bool = True) -> bool:
             logger.debug(f'Image {image_name} does not exist.')
 
         return result['exists']
+
+    def close(self):
+        self._closed = True

openhands/runtime/impl/remote/remote_runtime.py

@@ -13,6 +13,7 @@
     AgentRuntimeNotReadyError,
     AgentRuntimeUnavailableError,
 )
+from openhands.core.logger import openhands_logger as logger
 from openhands.events import EventStream
 from openhands.runtime.builder.remote import RemoteRuntimeBuilder
 from openhands.runtime.impl.action_execution.action_execution_client import (
@@ -66,12 +67,15 @@ def __init__(
             )
 
         self.runtime_builder = RemoteRuntimeBuilder(
-            self.config.sandbox.remote_runtime_api_url, self.config.sandbox.api_key
+            self.config.sandbox.remote_runtime_api_url,
+            self.config.sandbox.api_key,
+            self.session,
         )
         self.runtime_id: str | None = None
         self.runtime_url: str | None = None
         self.available_hosts: dict[str, int] = {}
         self._runtime_initialized: bool = False
+        self._closed: bool = False
 
     def _get_action_execution_server_host(self):
         return self.runtime_url
@@ -285,7 +289,8 @@ def _wait_until_alive(self):
             stop=tenacity.stop_after_delay(
                 self.config.sandbox.remote_runtime_init_timeout
             )
-            | stop_if_should_exit(),
+            | stop_if_should_exit()
+            | self.is_closed,
             reraise=True,
             retry=tenacity.retry_if_exception_type(AgentRuntimeNotReadyError),
             wait=tenacity.wait_fixed(2),
@@ -349,7 +354,11 @@ def _wait_until_alive_impl(self):
         raise AgentRuntimeNotReadyError()
 
     def close(self):
+        if self._closed:
+            return
+        self._closed = True
         if self.config.sandbox.keep_runtime_alive or self.attach_to_existing:
+            self.runtime_builder.close()
             super().close()
             return
         try:
@@ -359,11 +368,15 @@ def close(self):
                 json={'runtime_id': self.runtime_id},
             ):
                 self.log('debug', 'Runtime stopped.')
-        except Exception as e:
-            raise e
+        except Exception:
+            logger.warning('failed to stop remote runtime.')
         finally:
+            self.runtime_builder.close()
             super().close()
 
+    def is_closed(self):
+        return self._closed
+
     def _send_runtime_api_request(self, method, url, **kwargs):
         return send_request(self.session, method, url, **kwargs)
 

openhands/server/session/agent_session.py

@@ -16,11 +16,7 @@
 from openhands.runtime.base import Runtime
 from openhands.security import SecurityAnalyzer, options
 from openhands.storage.files import FileStore
-from openhands.utils.async_utils import call_async_from_sync, call_sync_from_async
-from openhands.utils.shutdown_listener import should_continue
-
-WAIT_TIME_BEFORE_CLOSE = 300
-WAIT_TIME_BEFORE_CLOSE_INTERVAL = 5
+from openhands.utils.async_utils import call_sync_from_async
 
 
 class AgentSession:
@@ -36,7 +32,7 @@ class AgentSession:
     controller: AgentController | None = None
     runtime: Runtime | None = None
     security_analyzer: SecurityAnalyzer | None = None
-    _initializing: bool = False
+    # _initializing: bool = False
     _closed: bool = False
     loop: asyncio.AbstractEventLoop | None = None
 
@@ -88,7 +84,6 @@ async def start(
         if self._closed:
             logger.warning('Session closed before starting')
             return
-        self._initializing = True
         self._create_security_analyzer(config.security.security_analyzer)
         await self._create_runtime(
             runtime_name=runtime_name,
@@ -111,36 +106,24 @@ async def start(
         )
         self._initializing = False
 
-    def close(self):
+    async def close(self):
         """Closes the Agent session"""
-        if self._closed:
-            return
-        self._closed = True
-        call_async_from_sync(self._close)
-
-    async def _close(self):
-        seconds_waited = 0
-        while self._initializing and should_continue():
-            logger.debug(
-                f'Waiting for initialization to finish before closing session {self.sid}'
-            )
-            await asyncio.sleep(WAIT_TIME_BEFORE_CLOSE_INTERVAL)
-            seconds_waited += WAIT_TIME_BEFORE_CLOSE_INTERVAL
-            if seconds_waited > WAIT_TIME_BEFORE_CLOSE:
-                logger.error(
-                    f'Waited too long for initialization to finish before closing session {self.sid}'
-                )
-                break
-        if self.event_stream is not None:
-            self.event_stream.close()
-        if self.controller is not None:
-            end_state = self.controller.get_state()
-            end_state.save_to_session(self.sid, self.file_store)
-            await self.controller.close()
-        if self.runtime is not None:
-            self.runtime.close()
-        if self.security_analyzer is not None:
-            await self.security_analyzer.close()
+        try:
+            if self._closed:
+                return
+            self._closed = True
+            if self.event_stream is not None:
+                self.event_stream.close()
+            if self.controller is not None:
+                end_state = self.controller.get_state()
+                end_state.save_to_session(self.sid, self.file_store)
+                await self.controller.close()
+            if self.runtime is not None:
+                self.runtime.close()
+            if self.security_analyzer is not None:
+                await self.security_analyzer.close()
+        except Exception:
+            logger.warning('error_closing_session', exc_info=True, stack_info=True)
 
     async def stop_agent_loop_for_error(self):
         if self.controller is not None:

openhands/server/session/session.py

@@ -72,7 +72,7 @@ async def close(self):
                 to=ROOM_KEY.format(sid=self.sid),
             )
         self.is_alive = False
-        self.agent_session.close()
+        asyncio.create_task(self.agent_session.close())
 
     async def initialize_agent(
         self,
@@ -124,9 +124,9 @@ async def initialize_agent(
                 selected_repository=selected_repository,
             )
         except Exception as e:
-            logger.exception(f'Error creating controller: {e}')
+            logger.exception(f'Error creating agent_session: {e}')
             await self.send_error(
-                f'Error creating controller. Please check Docker is running and visit `{TROUBLESHOOTING_URL}` for more debugging information..'
+                f'Error creating agent_session. Please check Docker is running and visit `{TROUBLESHOOTING_URL}` for more debugging information..'
             )
             return