Fix connection lost bug