Bug fixes for UI code