Refactoring websocket to have a nicer interface