More UI work & error handling