Fix UI update