Taking notes, taking notes, taking notes. Notepad as an external memory is just an amazing game changer for this. Doesn't even need to be deep, just a running tally of what you're thinking to get it to somewhere that you can reference it instead of running it on a loop in your short term memory. Makes fading in and out of it so much easier. I treat it like a conversation with myself half the time and the other half the time like I'm writing a reference book for someone else. Its worked well for fresh code, code to be, and analyzing existing code made by someone else. Cannot recommend it enough.
Note taking is great not just for recovering from distractions but also preventing them in the first place. Just jot down branching tasks that would've been a distraction and come back to them later instead of being the Hal fixing a lightbulb meme.