Small state machine changes