diff --git a/openhands/server/session/manager.py b/openhands/server/session/manager.py index 7b38c3b73828..8a5081baef35 100644 --- a/openhands/server/session/manager.py +++ b/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: diff --git a/tests/unit/test_manager.py b/tests/unit/test_manager.py index 7b16915df8f4..ecc08a37e0fd 100644 --- a/tests/unit/test_manager.py +++ b/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'