WebSocket refactor