Fix for docker leak (All-Hands-AI#4560)

openhands/runtime/impl/eventstream/eventstream_runtime.py

@@ -332,7 +332,7 @@ def _init_container(
                 f'Error: Instance {self.container_name} FAILED to start container!\n'
             )
             logger.exception(e)
-            self.close(close_client=False)
+            self.close()
             raise e
 
     def _attach_to_container(self):
@@ -393,11 +393,10 @@ def _wait_until_alive(self):
             logger.error(msg)
             raise RuntimeError(msg)
 
-    def close(self, close_client: bool = True, rm_all_containers: bool = True):
+    def close(self, rm_all_containers: bool = True):
         """Closes the EventStreamRuntime and associated objects
 
         Parameters:
-        - close_client (bool): Whether to close the DockerClient
         - rm_all_containers (bool): Whether to remove all containers with the 'openhands-sandbox-' prefix
         """
 
@@ -407,6 +406,9 @@ def close(self, close_client: bool = True, rm_all_containers: bool = True):
         if self.session:
             self.session.close()
 
+        if self.attach_to_existing:
+            return
+
         try:
             containers = self.docker_client.containers.list(all=True)
             for container in containers:
@@ -431,9 +433,6 @@ def close(self, close_client: bool = True, rm_all_containers: bool = True):
         except docker.errors.NotFound:  # yes, this can happen!
             pass
 
-        if close_client:
-            self.docker_client.close()
-
     def run_action(self, action: Action) -> Observation:
         if isinstance(action, FileEditAction):
             return self.edit(action)

openhands/runtime/impl/remote/remote_runtime.py

@@ -318,7 +318,7 @@ def _wait_until_alive(self):
             raise RuntimeError(msg)
 
     def close(self, timeout: int = 10):
-        if self.config.sandbox.keep_remote_runtime_alive:
+        if self.config.sandbox.keep_remote_runtime_alive or self.attach_to_existing:
             self.session.close()
             return
         if self.runtime_id:

openhands/server/listen.py

@@ -247,8 +247,10 @@ async def attach_session(request: Request, call_next):
             status_code=status.HTTP_404_NOT_FOUND,
             content={'error': 'Session not found'},
         )
-
-    response = await call_next(request)
+    try:
+        response = await call_next(request)
+    finally:
+        await session_manager.detach_from_conversation(request.state.conversation)
     return response
 
 

openhands/server/session/conversation.py

@@ -1,9 +1,12 @@
+import asyncio
+
 from openhands.core.config import AppConfig
 from openhands.events.stream import EventStream
 from openhands.runtime import get_runtime_cls
 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_sync_from_async
 
 
 class Conversation:
@@ -37,3 +40,6 @@ def __init__(
 
     async def connect(self):
         await self.runtime.connect()
+
+    async def disconnect(self):
+        asyncio.create_task(call_sync_from_async(self.runtime.close))

openhands/server/session/manager.py

@@ -53,6 +53,9 @@ async def attach_to_conversation(self, sid: str) -> Conversation | None:
         await c.connect()
         return c
 
+    async def detach_from_conversation(self, conversation: Conversation):
+        await conversation.disconnect()
+
     async def send(self, sid: str, data: dict[str, object]) -> bool:
         """Sends data to the client."""
         session = self.get_session(sid)