From c3da25febc07d0b696aa08844dabb1e9673e7436 Mon Sep 17 00:00:00 2001
From: tofarr <tofarr@gmail.com>
Date: Fri, 25 Oct 2024 09:53:39 -0600
Subject: [PATCH] Fix for docker leak (#4560)

---
 .../runtime/impl/eventstream/eventstream_runtime.py   | 11 +++++------
 openhands/runtime/impl/remote/remote_runtime.py       |  2 +-
 openhands/server/listen.py                            |  6 ++++--
 openhands/server/session/conversation.py              |  6 ++++++
 openhands/server/session/manager.py                   |  3 +++
 5 files changed, 19 insertions(+), 9 deletions(-)

diff --git a/openhands/runtime/impl/eventstream/eventstream_runtime.py b/openhands/runtime/impl/eventstream/eventstream_runtime.py
index 062fda2488..1cf4096922 100644
--- a/openhands/runtime/impl/eventstream/eventstream_runtime.py
+++ b/openhands/runtime/impl/eventstream/eventstream_runtime.py
@@ -332,7 +332,7 @@ class EventStreamRuntime(Runtime):
                 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 @@ class EventStreamRuntime(Runtime):
             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 @@ class EventStreamRuntime(Runtime):
         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 @@ class EventStreamRuntime(Runtime):
         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)
diff --git a/openhands/runtime/impl/remote/remote_runtime.py b/openhands/runtime/impl/remote/remote_runtime.py
index b3d60488c8..8c25401af9 100644
--- a/openhands/runtime/impl/remote/remote_runtime.py
+++ b/openhands/runtime/impl/remote/remote_runtime.py
@@ -318,7 +318,7 @@ class RemoteRuntime(Runtime):
             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:
diff --git a/openhands/server/listen.py b/openhands/server/listen.py
index a556d78adc..9647b7bc84 100644
--- a/openhands/server/listen.py
+++ b/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
 
 
diff --git a/openhands/server/session/conversation.py b/openhands/server/session/conversation.py
index 8e99864635..ad880840e5 100644
--- a/openhands/server/session/conversation.py
+++ b/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 @@ class Conversation:
 
     async def connect(self):
         await self.runtime.connect()
+
+    async def disconnect(self):
+        asyncio.create_task(call_sync_from_async(self.runtime.close))
diff --git a/openhands/server/session/manager.py b/openhands/server/session/manager.py
index 3ce0e23235..6c650feee7 100644
--- a/openhands/server/session/manager.py
+++ b/openhands/server/session/manager.py
@@ -53,6 +53,9 @@ class SessionManager:
         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)