Fix issues with disconnecting