Answering my own question. This was trivially simple in my code. My code was doing deep sleep in the setup() function. When waking up from deep-sleep the setup function is called again. But if the device is powered by USB then deep sleep fails and we will reach the loop() function. So any special code that I want for USB power should simply be put in loop().
D