Using a State Management Library

In this guide we are explaining how you could use React Flow with the state management library Zustand (opens in a new tab). We will build a little app where every node has a color chooser that updates its background color. In this guide we are are using Zustand, because we are already using it internally for React Flow, but of course you can use any other library like Redux (opens in a new tab), Recoil (opens in a new tab) or Jotai (opens in a new tab) as well.

As you might have seen in the previous guides and examples, React Flow can easily be used with a local component state for handling the nodes and edges of your diagram. When your app grows and you want to alter your state from within your nodes for example, things can get more complex. To avoid passing down functions through the node data field, you could use a React context (opens in a new tab) or add a state management library as explained in this guide.

Install Zustand

As mentioned above we are using Zustand in this example. Zustand is a bit like Redux: you have a central store with actions to alter your state and hooks to access your state. You can install Zustand via:

npm install --save zustand

Create a Store

Zustand lets you create a hook for accessing the values and functions of your store. We put the nodes and edges and the onNodesChange, onEdgesChange, onConnect, setNodes and setEdges functions in the store to get the basic interactivity for our graph:

export default function App() {
  const data: string = "world"

  return <h1>Hello {data}</h1>
}

That's the basic setup. We now have a store with nodes and edges that can handle the changes (dragging, selecting or removing a node or edge) triggered by React Flow. When you take a look at the App.tsx file, you can see that it's kept nice and clean. All the data and actions are now part of the store and can be accessed with the useStore hook.

Implement a Color Change Action

We add a new updateNodeColor action to update the data.color field of a specific node. For this we pass the node id and the new color to the action, iterate over the nodes and update the matching one with the new color:

updateNodeColor: (nodeId: string, color: string) => {
  set({
    nodes: get().nodes.map((node) => {
      if (node.id === nodeId) {
        // it's important to create a new object here, to inform React Flow about the changes
        return { ...node, data: { ...node.data, color } };
      }
 
      return node;
    }),
  });
};

This new action can now be used in a React component like this:

const updateNodeColor = useStore((s) => s.updateNodeColor);
...
<button onClick={() => updateNodeColor(nodeId, color)} />;

Add a Color Chooser Node

In this step we implement the ColorChooserNode component and call the updateNodeColor when the user changes the color. The custom part of the color chooser node is the color input.

<input
  type="color"
  defaultValue={data.color}
  onChange={(evt) => updateNodeColor(id, evt.target.value)}
  className="nodrag"
/>

We add the nodrag class name so that the user doesn't drag the node by mistake when changing the color and call the updateNodeColor in the onChange event handler.

export default function App() {
  const data: string = "world"

  return <h1>Hello {data}</h1>
}

You can now click on a color chooser and change the background of a node.