await closing the controller (#1751)

* await closing the controller

* Update manager.py

* Cleanly exit

* Update agent.py

---------

Co-authored-by: Yufan Song <33971064+yufansong@users.noreply.github.com>
Co-authored-by: Jim Su <jimsu@protonmail.com>

opendevin/server/agent/agent.py

@@ -138,6 +138,10 @@ class AgentUnit:
             await self.send(event.to_dict())
         return
 
-    def close(self):
+    async def close(self):
+        """Cleanly exits an AgentUnit.
+
+        Execution awaited by the AgentManager.
+        """
         if self.controller is not None:
-            self.controller.close()
+            await self.controller.close()

opendevin/server/agent/manager.py

@@ -1,4 +1,4 @@
-import atexit
+import asyncio, atexit
 
 from opendevin.core.logger import opendevin_logger as logger
 from opendevin.server.session import session_manager
@@ -35,6 +35,14 @@ class AgentManager:
         await self.sid_to_agent[sid].dispatch(action, data)
 
     def close(self):
+        try:
+            loop = asyncio.get_event_loop()
+        except RuntimeError:
+            loop = asyncio.new_event_loop()
+            asyncio.set_event_loop(loop)
+        loop.run_until_complete(self._close())
+
+    async def _close(self):
         logger.info(f'Closing {len(self.sid_to_agent)} agent(s)...')
         for sid, agent in self.sid_to_agent.items():
-            agent.close()
+            await agent.close()