refactor to fix state error with websocket