Fix messages before disconnect problem