K

Making a state machine in TypeScript

Learn to code your own state machine in TypeScript to cleanly handle application state.

Posted: November 25, 2022
Find more posts about:

State machines are a nifty programming algorithm to control application flow. In it’s simplest form, a state machine has a set of states that it can exist in and a set of actions that can occur in each of those states.

A great example is a traffic light. In a standard traffic light, the three states are green, yellow, and red. There is one action, change. When the light is in the green state and receives the change action, it will always go into the yellow state. There is a clear order that can’t be violated.

A state machine for a traffic light

The state machine knows the flow of these states and their actions, and it will never allow your program to get into an incorrect state (as long as your states and actions are defined correctly!). This gives you much cleaner and predictable code; instead of needing to check a bunch of conditionals like if (!error && !loading && response), you can just check the state using if (state === 'ready').

If you plan on using many state machines that could have nested logic, I recommend reaching for a strong library like xstate. But if you just need a one-off, simple machine, you can make one yourself!

For this example, we’ll make a state machine to run a robotic fortune-teller. Think of this like one of those Zoltar fortune teller machines you might see at an old carnival; the user can push a button, and the robot will tell them either a good or bad fortune.

An AI-generated robot fortune-teller

Though this is a silly example, think about how this model can be applied in real applications. For example, instead of getting back a good or bad fortune from a robot, you could get back a response or an error from an API!

Modeling our states

First, we’ll need to clearly define the states and actions. It’s a great help to draw this out visually.

A sketch of the flow for our fortune-teller

The robot will first greet the user and invite them to push the button to receive a fortune. Whenever the user pushes the button, the robot will take a moment to look into his crystal ball to find a fortune. Once the robot receives the fortune back, it will display the fortune to the user. Since it could be a good or bad fortune, we can display those differently! The bad fortune will flash scary lights and play spooky music, while a good fortune will play happy music. Regardless of the type of fortune the user gets back, they can push the button again to get another fortune.

Now that we have our robot behavior drawn out, let’s rewrite our states and actions to be more program-focused instead of behavior focused.

The same robot flow, but with program-like states and actions

We first show our greet state. When the user pushes the button, we’ll send a request to whatever service or man-behind-the-curtain we’re using to generate the fortunes, and we’ll show our waiting state. Once we get a response back, we can show the fortune based on whether it’s good or bad. Once the response is shown, the another response can be sent, which puts our robot back in the waiting state.

With our states and actions clearly defined, we can finally start writing some code! While we could manage all of this with strings, we’ll be using enums to help prevent any accidental mispellings or errors.

enum MachineStates {
	greet = 'greet',
	waiting = 'waiting',
	show_bad = 'show_bad',
	show_good = 'show_good',
}

enum MachineActions {
	send_request = 'send_request',
	receive_bad = 'receive_bad',
	receive_good = 'receive_good',
}

type MachineState = keyof typeof MachineStates
type MachineAction = keyof typeof MachineActions

We’ve also created two types based on our enums. Using the keyof typeof keywords lets us pull all the keys of the enums, which is essentially equivalent to the following:

type MachineState = 'greet' | 'waiting' | 'show_bad' | 'show_good'
type MachineAction = 'send_request' | 'receive_bad' | 'receive_good'

However, we get the added benefit of these types automatically updating any time we alter our enums.

Building the state machine

Since we have our actions and states defined, we need to lay out the same model that we had in our diagram. Our machine needs to know which actions are allowed from which states, and what to do when it receives a certain action. This is just going to be a JSON object, so let’s mock out what it might look like so that we can then define a type for it.

const stateMachine = {
  'greet': {
    'send_request': 'waiting'
  },
  'waiting': {
    'receive_bad': 'show_bad',
    'receive_good': 'show_good'
  }
  // ...
}

We can see that each key in our state machine needs to be one of our possible states. We’ll look at the first state, greet, as an example. When we’re in the greet state, the only action we can take is send_request. When that action occurs, it takes us to the waiting state. So we have the top level keys being possible states, second level keys being possible actions, and our values being the resulting state. Let’s make a type for this.

type StateMachine = {
	[state in MachineStates]: {
		[action in MachineActions]?: MachineState
	}
}

For our top level keys, we’re essentially looping through all of the states in our MachineStates enum. That object’s keys come from a loop of all of our MachineActions, which point to a single resulting MachineState.

Nota bene: we’ve included a ? after the machine actions. This means that all of our states are required to be filled out, but each state doesn’t have to accept every action.

With our type defined, let’s fill out our actual machine.

const stateMachine: StateMachine = {
  [MachineStates.greet]: {
    [MachineActions.send_request]: MachineStates.waiting
  },
  [MachineStates.waiting]: {
    [MachineActions.receive_bad]: MachineStates.show_bad,
    [MachineActions.receive_good]: MachineStates.show_good,
  },
  [MachineStates.show_bad]: {
    [MachineActions.send_request]: MachineStates.waiting,
  },
  [MachineStates.show_good]: {
    [MachineActions.send_request]: MachineStates.waiting,
  }
}

Great! Now that we have our machine modeled and built, we need a way to utilize the machine.

Querying the state machine

We’re going to make a function that accepts the current state and the action that occurs then uses the state machine to tell us the new state we should be in.

const getNextState = (
	state: MachineState,
	action: MachineAction
) => {
  const newState = stateMachine[state][action];
  return newState;
}

Our function just has to lookup the new state we should be in using our JSON object! We can even have it throw an error if we try to perform an action that’s impossible from our current state.

const getNextState = (
	state: MachineState,
	action: MachineAction
) => {
  const newState = stateMachine[state][action];
    if (!newState) {
    throw new Error(`Action ${action} is not a valid step from state ${state}`);
  }
  return newState;
}

One of the nicest things about this machine is that it’s stateless. The machine isn’t keeping track of what state we’re in; it just tells you how to respond based on a certain state and action. This means that the state machine will be really easy to test 😎

Using the state machine in an application

Fantastic! Now that we have our machine built and know how to use it, we can plug it into an application. Instead of actually controlling a robot, though, it’ll be easier to demonstrate how it would be used in a React component.

We’ll start off with a barebones functional component.

const FortuneTeller = () => {
	return (
		<div>
			<h1>Fortune Teller</h1>
		</div>
	)
}

Since our state machine is stateless, we’ll need to keep track of our state in this component using the useState hook. We can initialize the state to the greet state. We’ll also eventually need a variable for our fortune, so we’ll create it now too.

const FortuneTeller = () => {
  const [state, setState] = useState<MachineState>(MachineStates.greet)
  const [fortune, setFortune] = useState<string>('');
  return (
		<div>
			<h1>Fortune Teller</h1>
		</div>
	)
}

And now, one of the coolest parts about using a state machine: since we have our states defined, we know all the states that our frontend UI can be in! We can build them out without needing to worry about complicated conditionals.

const FortuneTeller = () => {
  const [state, setState] = useState<MachineState>(MachineStates.greet)
  return (
		<div>
			<h1>Fortune Teller</h1>
      {state === MachineStates.greet && (
        <p>Press the button to get your fortune!</p>
	  )}
      {state === MachineStates.waiting && (
	    <p>Waiting for fortune...</p>
	  )}
      {state === MachineStates.show_bad && (
        <div>
          <p>Uh oh! You got a bad fortune!</p>
          <p>{fortune}</p>
        </div>
      )}
      {state === MachineStates.show_good && (
        <div>
          <p>Nice! You got a good fortune!</p>
          <p>{fortune}</p>
        </div>
      )}
		</div>
	)
}

We just add a check to see which state we’re in, and display the corresponding UI. Next, we can add our button in so the user can interact and request a fortune.

const FortuneTeller = () => {
  const [state, setState] = useState<MachineState>(MachineStates.greet)
  return (
		<div>
			<h1>Fortune Teller</h1>
      {state === MachineStates.greet && <p>Press the button to get your fortune!</p>}
      {state === MachineStates.waiting && <p>Waiting for fortune...</p>}
      {state === MachineStates.show_bad && (
        <div>
          <p>Uh oh! You got a bad fortune!</p>
          <p>{fortune}</p>
        </div>
      )}
      {state === MachineStates.show_good && (
        <div>
          <p>Nice! You got a good fortune!</p>
          <p>{fortune}</p>
        </div>
      )}
      <button
        onClick={() => requestFortune()}
      >
        Get a fortune!
      </button>
		</div>
	)
}

We need to be careful, though. What happens if the user clicks to get a fortune, but then clicks again while we’re in the waiting state? We know from our state machine that the state waiting can’t accept the action send_request!

We have two options:

  • don’t display the button if we’re in the waiting state
  • disable the button if we’re in the waiting state

Both of these options are valid! You can choose which fits best with your project.

{state !== MachineStates.waiting && (
  <button
    onClick={() => requestFortune()}
  >
	Get a fortune!
  </button>
)}
// OR
<button
  onClick={() => requestFortune()}
  disabled={state === MachineStates.waiting}
>
  Get a fortune!
</button>

Great! Now the last step is to implement the function that actually requests a fortune, requestFortune():

const requestFortune = async (): Promise<void> => {
    setState(state => getNextState(state, MachineActions.send_request))
    const [isGoodFortune, fortuneResponse] = await seeFuture();
    setState(state => getNextState(
      state,
      isGoodFortune ?
        MachineActions.receive_good :
        MachineActions.receive_bad
      )
    )
    setFortune(fortuneResponse)
  }

When we begin this function call, we pass in our current state and the action that we’re taking, send_request. Look at that! We don’t have to worry about what state we want to go to because the state machine takes care of that for us. We just have to tell it what state we’re in and what action we want to do.

We make our call to our fortune service seeFuture() and await the response (hey, seeing the future isn’t easy!), then send off our next action depending on if the fortune was good or bad. Finally, we update the fortune so it can be displayed to the user.

Nota bene: whenever we update the state using the useState setter function, we have to use the function version. Otherwise, React might try to batch your updates together and skip a step in your machine, which will cause your machine to throw an error.

// ✅ Do this
setState(state => getNewState(state, MachineActions.send_request))

// 🚫 Don't do this!
setState(getNewState(state, MachineActions.send_request))

Go forth and machine!

And there we go! You now have a homebrew state machine to manage your application and write more predictable, maintainable code. Have fun experimenting!