Clean up towards use of state machine