Alright, so it turns out that I originally implemented this in a way more robust way than I remembered (surprising!) and so these issues have turned out to be straightforward to fix. I was very worried they were going to be really tricky.