Add TODO for cleanup