Clean up controller logic