State restoration fixes