Fixing up bugs from UI overhaul