The planning, which can also be done in concert with AI, before getting the AI to write the code helps with not losing track of what you're building. In my current workflow, I never tell AI to just code X, I always approach it in such a way to leave space for research and breaking it up into multiple units of work (tickets) if the scope of the feature or idea calls for it.