Lockup Resiliency and Asyncio Improvements (#4221)

openhands/server/listen.py

@@ -430,7 +430,9 @@ async def list_files(request: Request, path: str | None = None):
             content={'error': 'Runtime not yet initialized'},
         )
     runtime: Runtime = request.state.session.agent_session.runtime
-    file_list = runtime.list_files(path)
+    file_list = await asyncio.get_event_loop().run_in_executor(
+        None, runtime.list_files, path
+    )
     if path:
         file_list = [os.path.join(path, f) for f in file_list]
 
@@ -451,6 +453,7 @@ async def list_files(request: Request, path: str | None = None):
         return file_list
 
     file_list = filter_for_gitignore(file_list, '')
+
     return file_list
 
 
@@ -478,7 +481,7 @@ async def select_file(file: str, request: Request):
 
     file = os.path.join(runtime.config.workspace_mount_path_in_sandbox, file)
     read_action = FileReadAction(file)
-    observation = runtime.run_action(read_action)
+    observation = await runtime.async_run_action(read_action)
 
     if isinstance(observation, FileReadObservation):
         content = observation.content
@@ -720,7 +723,7 @@ async def save_file(request: Request):
             runtime.config.workspace_mount_path_in_sandbox, file_path
         )
         write_action = FileWriteAction(file_path, content)
-        observation = runtime.run_action(write_action)
+        observation = await runtime.async_run_action(write_action)
 
         if isinstance(observation, FileWriteObservation):
             return JSONResponse(

openhands/server/session/agent_session.py

@@ -1,6 +1,4 @@
 import asyncio
-import concurrent.futures
-from threading import Thread
 from typing import Callable, Optional
 
 from openhands.controller import AgentController
@@ -32,7 +30,7 @@ class AgentSession:
     runtime: Runtime | None = None
     security_analyzer: SecurityAnalyzer | None = None
     _closed: bool = False
-    loop: asyncio.AbstractEventLoop
+    loop: asyncio.AbstractEventLoop | None = None
 
     def __init__(self, sid: str, file_store: FileStore):
         """Initializes a new instance of the Session class
@@ -45,7 +43,6 @@ class AgentSession:
         self.sid = sid
         self.event_stream = EventStream(sid, file_store)
         self.file_store = file_store
-        self.loop = asyncio.new_event_loop()
 
     async def start(
         self,
@@ -73,17 +70,9 @@ class AgentSession:
                 'Session already started. You need to close this session and start a new one.'
             )
 
-        self.thread = Thread(target=self._run, daemon=True)
-        self.thread.start()
-
-        def coro_callback(task):
-            fut: concurrent.futures.Future = concurrent.futures.Future()
-            try:
-                fut.set_result(task.result())
-            except Exception as e:
-                logger.error(f'Error starting session: {e}')
-
-        coro = self._start(
+        asyncio.get_event_loop().run_in_executor(
+            None,
+            self._start_thread,
             runtime_name,
             config,
             agent,
@@ -93,9 +82,12 @@ class AgentSession:
             agent_configs,
             status_message_callback,
         )
-        asyncio.run_coroutine_threadsafe(coro, self.loop).add_done_callback(
-            coro_callback
-        )  # type: ignore
+
+    def _start_thread(self, *args):
+        try:
+            asyncio.run(self._start(*args), debug=True)
+        except RuntimeError:
+            logger.info('Session Finished')
 
     async def _start(
         self,
@@ -108,6 +100,7 @@ class AgentSession:
         agent_configs: dict[str, AgentConfig] | None = None,
         status_message_callback: Optional[Callable] = None,
     ):
+        self.loop = asyncio.get_running_loop()
         self._create_security_analyzer(config.security.security_analyzer)
         self._create_runtime(runtime_name, config, agent, status_message_callback)
         self._create_controller(
@@ -125,10 +118,6 @@ class AgentSession:
             self.controller.agent_task = self.controller.start_step_loop()
             await self.controller.agent_task  # type: ignore
 
-    def _run(self):
-        asyncio.set_event_loop(self.loop)
-        self.loop.run_forever()
-
     async def close(self):
         """Closes the Agent session"""
 
@@ -143,10 +132,8 @@ class AgentSession:
         if self.security_analyzer is not None:
             await self.security_analyzer.close()
 
-        self.loop.call_soon_threadsafe(self.loop.stop)
-        if self.thread:
-            # We may be closing an agent_session that was never actually started
-            self.thread.join()
+        if self.loop:
+            self.loop.call_soon_threadsafe(self.loop.stop)
 
         self._closed = True
 

openhands/server/session/session.py

@@ -162,9 +162,10 @@ class Session:
                         'Model does not support image upload, change to a different model or try without an image.'
                     )
                     return
-        asyncio.run_coroutine_threadsafe(
-            self._add_event(event, EventSource.USER), self.agent_session.loop
-        )  # type: ignore
+        if self.agent_session.loop:
+            asyncio.run_coroutine_threadsafe(
+                self._add_event(event, EventSource.USER), self.agent_session.loop
+            )  # type: ignore
 
     async def _add_event(self, event, event_source):
         self.agent_session.event_stream.add_event(event, EventSource.USER)