Add proper websocket close handling, fix on_close callback