Refactor state machines after recent changes