Fix state cleanup