Refactor task controller