From 9515ac5e6231f88f907bfedfb8fcb567d9cc67f7 Mon Sep 17 00:00:00 2001
From: tofarr <tofarr@gmail.com>
Date: Mon, 6 Jan 2025 14:26:48 -0700
Subject: [PATCH] Feat - browser client can now close sessions. (#6088)

---
 .../server/routes/manage_conversations.py     |  2 +-
 openhands/server/session/manager.py           | 22 ++++++++++++++++---
 tests/unit/test_manager.py                    |  2 +-
 3 files changed, 21 insertions(+), 5 deletions(-)

diff --git a/openhands/server/routes/manage_conversations.py b/openhands/server/routes/manage_conversations.py
index a48c7286879a..235b5801f24d 100644
--- a/openhands/server/routes/manage_conversations.py
+++ b/openhands/server/routes/manage_conversations.py
@@ -168,7 +168,7 @@ async def delete_conversation(
         return False
     is_running = await session_manager.is_agent_loop_running(conversation_id)
     if is_running:
-        return False
+        await session_manager.close_session(conversation_id)
     await conversation_store.delete_metadata(conversation_id)
     return True
 
diff --git a/openhands/server/session/manager.py b/openhands/server/session/manager.py
index 60b5bd2675af..cc08d87466e1 100644
--- a/openhands/server/session/manager.py
+++ b/openhands/server/session/manager.py
@@ -156,6 +156,10 @@ async def _process_message(self, message: dict):
             flag = self._has_remote_connections_flags.get(sid)
             if flag:
                 flag.set()
+        elif message_type == 'close_session':
+            sid = data['sid']
+            if sid in self._local_agent_loops_by_sid:
+                await self._on_close_session(sid)
         elif message_type == 'session_closing':
             # Session closing event - We only get this in the event of graceful shutdown,
             # which can't be guaranteed - nodes can simply vanish unexpectedly!
@@ -419,7 +423,7 @@ async def disconnect_from_session(self, connection_id: str):
         if should_continue():
             asyncio.create_task(self._cleanup_session_later(sid))
         else:
-            await self._close_session(sid)
+            await self._on_close_session(sid)
 
     async def _cleanup_session_later(self, sid: str):
         # Once there have been no connections to a session for a reasonable period, we close it
@@ -451,10 +455,22 @@ async def _cleanup_session(self, sid: str) -> bool:
                 json.dumps({'sid': sid, 'message_type': 'session_closing'}),
             )
 
-        await self._close_session(sid)
+        await self._on_close_session(sid)
         return True
 
-    async def _close_session(self, sid: str):
+    async def close_session(self, sid: str):
+        session = self._local_agent_loops_by_sid.get(sid)
+        if session:
+            await self._on_close_session(sid)
+
+        redis_client = self._get_redis_client()
+        if redis_client:
+            await redis_client.publish(
+                'oh_event',
+                json.dumps({'sid': sid, 'message_type': 'close_session'}),
+            )
+
+    async def _on_close_session(self, sid: str):
         logger.info(f'_close_session:{sid}')
 
         # Clear up local variables
diff --git a/tests/unit/test_manager.py b/tests/unit/test_manager.py
index c2a61104a864..144f79f9f491 100644
--- a/tests/unit/test_manager.py
+++ b/tests/unit/test_manager.py
@@ -286,7 +286,7 @@ async def test_cleanup_session_connections():
                 }
             )
 
-            await session_manager._close_session('session1')
+            await session_manager._on_close_session('session1')
 
             remaining_connections = session_manager.local_connection_id_to_session_id
             assert 'conn1' not in remaining_connections