Improve error handling on controller