Fix connection check (All-Hands-AI#5787)

openhands/server/session/manager.py

@@ -421,8 +421,8 @@ async def _close_session(self, sid: str):
         # Clear up local variables
         connection_ids_to_remove = list(
             connection_id
-            for connection_id, sid in self.local_connection_id_to_session_id.items()
-            if sid == sid
+            for connection_id, conn_sid in self.local_connection_id_to_session_id.items()
+            if sid == conn_sid
         )
         logger.info(f'removing connections: {connection_ids_to_remove}')
         for connnnection_id in connection_ids_to_remove:

tests/unit/test_manager.py

@@ -237,3 +237,36 @@ async def test_add_to_cluster_event_stream():
         'oh_event',
         '{"sid": "new-session-id", "message_type": "event", "data": {"event_type": "some_event"}}',
     )
+
+
+@pytest.mark.asyncio
+async def test_cleanup_session_connections():
+    sio = get_mock_sio()
+    with (
+        patch('openhands.server.session.manager._REDIS_POLL_TIMEOUT', 0.01),
+        patch(
+            'openhands.server.session.manager.SessionManager._redis_subscribe',
+            AsyncMock(),
+        ),
+    ):
+        async with SessionManager(
+            sio, AppConfig(), InMemoryFileStore()
+        ) as session_manager:
+            session_manager.local_connection_id_to_session_id.update(
+                {
+                    'conn1': 'session1',
+                    'conn2': 'session1',
+                    'conn3': 'session2',
+                    'conn4': 'session2',
+                }
+            )
+
+            await session_manager._close_session('session1')
+
+            remaining_connections = session_manager.local_connection_id_to_session_id
+            assert 'conn1' not in remaining_connections
+            assert 'conn2' not in remaining_connections
+            assert 'conn3' in remaining_connections
+            assert 'conn4' in remaining_connections
+            assert remaining_connections['conn3'] == 'session2'
+            assert remaining_connections['conn4'] == 'session2'