You know, I think you've pegged it. Keeping it in my head is probably the main factor.

I write down lots of thoughts and tiny tasks in a notebook as I'm working, and check off the tasks as I go, but that's not really setting up for context-switching. When I'm writing code, I'm following inspiration: playing out the idea in my head, or rather, using the act of writing code to cause the idea to become clarified and completed. While ramping up and while coding, the idea is taking shape.

Maintaining notes designed to enable context-switching sounds like something I could try. This is probably the key. This would enable working in smaller chunks of time, too.

Weird: Right at this moment, the thought terrifies me. I wonder if I'd have to completely change my creative process to make such notes.

Hmm.