If you're using AI like this, it seems to me that it would be perfectly reasonable to make a separate branch, allow for micro-commits, and squash once a "complete chunk of work" is done.

That said, having a single option that rewinds LLM context and code state is better than having to do both separately.