Add check before disconnecting