Using a State Management Library

For this guide we assume that you already know about the [core concepts](/learn/concepts/core-concepts) of React Flow and how to implement [custom nodes](/learn/customization/custom-nodes). You should also be familiar with the concepts of state management libraries and how to use them.

In this guide, we explain how to use React Flow with the state management library Zustand. We will build a small app where each node features a color chooser that updates its background color. We chose Zustand for this guide because React Flow already uses it internally, but you can easily use other state management libraries such as Redux, Recoil or Jotai

As demonstrated in previous guides and examples, React Flow can easily be used with a local component state to manage nodes and edges in your diagram. However, as your application grows and you need to update the state from within individual nodes, managing this state can become more complex. Instead of passing functions through the node’s data field, you can use a React context or integrate a state management library like Zustand, as outlined 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:

pnpm 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:

Example: learn/state-management

App.tsx
import { useShallow } from 'zustand/react/shallow';
import { ReactFlow } from '@xyflow/react';
 
import '@xyflow/react/dist/style.css';
 
import useStore from './store';
 
const selector = (state) => ({
  nodes: state.nodes,
  edges: state.edges,
  onNodesChange: state.onNodesChange,
  onEdgesChange: state.onEdgesChange,
  onConnect: state.onConnect,
});
 
function Flow() {
  const { nodes, edges, onNodesChange, onEdgesChange, onConnect } = useStore(
    useShallow(selector),
  );
 
  return (
    <ReactFlow
      nodes={nodes}
      edges={edges}
      onNodesChange={onNodesChange}
      onEdgesChange={onEdgesChange}
      onConnect={onConnect}
      fitView
      colorMode="system"
    />
  );
}
 
export default Flow;
edges.ts
import { type Edge } from '@xyflow/react';
 
export const initialEdges = [
  { id: 'e1-2', source: '1', target: '2' },
  { id: 'e2-3', source: '2', target: '3' },
] as Edge[];
index.css
html,
body {
  margin: 0;
  font-family: sans-serif;
}
 
#app {
  width: 100vw;
  height: 100vh;
}
index.html
<!doctype html>
<html lang="en">
  <head>
    <meta charset="UTF-8" />
    <meta name="viewport" content="width=device-width, initial-scale=1.0" />
    <title>React Flow Example</title>
  </head>
  <body>
    <div id="app"></div>
    <script type="module" src="./index.tsx"></script>
  </body>
</html>
index.tsx
import { createRoot } from 'react-dom/client';
import App from './App';
 
import './index.css';
 
const container = document.querySelector('#app');
const root = createRoot(container);
 
root.render(<App />);
nodes.ts
import { type AppNode } from './types';
 
export const initialNodes = [
  {
    id: '1',
    type: 'input',
    data: { label: 'Input' },
    position: { x: 250, y: 25 },
  },
 
  {
    id: '2',
    data: { label: 'Default' },
    position: { x: 100, y: 125 },
  },
  {
    id: '3',
    type: 'output',
    data: { label: 'Output' },
    position: { x: 250, y: 250 },
  },
] as AppNode[];
store.ts
import { create } from 'zustand';
import { addEdge, applyNodeChanges, applyEdgeChanges } from '@xyflow/react';
 
import { initialNodes } from './nodes';
import { initialEdges } from './edges';
import { type AppState } from './types';
 
// this is our useStore hook that we can use in our components to get parts of the store and call actions
const useStore = create<AppState>((set, get) => ({
  nodes: initialNodes,
  edges: initialEdges,
  onNodesChange: (changes) => {
    set({
      nodes: applyNodeChanges(changes, get().nodes),
    });
  },
  onEdgesChange: (changes) => {
    set({
      edges: applyEdgeChanges(changes, get().edges),
    });
  },
  onConnect: (connection) => {
    set({
      edges: addEdge(connection, get().edges),
    });
  },
  setNodes: (nodes) => {
    set({ nodes });
  },
  setEdges: (edges) => {
    set({ edges });
  },
}));
 
export default useStore;
types.ts
import {
  type Edge,
  type Node,
  type OnNodesChange,
  type OnEdgesChange,
  type OnConnect,
} from '@xyflow/react';
 
export type AppNode = Node;
 
export type AppState = {
  nodes: AppNode[];
  edges: Edge[];
  onNodesChange: OnNodesChange<AppNode>;
  onEdgesChange: OnEdgesChange;
  onConnect: OnConnect;
  setNodes: (nodes: AppNode[]) => void;
  setEdges: (edges: Edge[]) => void;
};

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.

Example: learn/state-management-2

App.tsx
import { useShallow } from 'zustand/react/shallow';
import { ReactFlow } from '@xyflow/react';
 
import '@xyflow/react/dist/style.css';
 
import useStore from './store';
import ColorChooserNode from './ColorChooserNode';
 
const nodeTypes = { colorChooser: ColorChooserNode };
 
const selector = (state) => ({
  nodes: state.nodes,
  edges: state.edges,
  onNodesChange: state.onNodesChange,
  onEdgesChange: state.onEdgesChange,
  onConnect: state.onConnect,
});
 
function Flow() {
  const { nodes, edges, onNodesChange, onEdgesChange, onConnect } = useStore(
    useShallow(selector),
  );
 
  return (
    <ReactFlow
      nodes={nodes}
      edges={edges}
      onNodesChange={onNodesChange}
      onEdgesChange={onEdgesChange}
      onConnect={onConnect}
      nodeTypes={nodeTypes}
      fitView
      colorMode="system"
    />
  );
}
 
export default Flow;
ColorChooserNode.tsx
import { Handle, type NodeProps, Position } from '@xyflow/react';
 
import useStore from './store';
import { type ColorNode } from './types';
 
function ColorChooserNode({ id, data }: NodeProps<ColorNode>) {
  const updateNodeColor = useStore((state) => state.updateNodeColor);
 
  return (
    <div style={{ backgroundColor: data.color, borderRadius: 10 }}>
      <Handle type="target" position={Position.Top} />
      <div style={{ padding: 20 }}>
        <input
          type="color"
          defaultValue={data.color}
          onChange={(evt) => updateNodeColor(id, evt.target.value)}
          className="nodrag"
        />
      </div>
      <Handle type="source" position={Position.Bottom} />
    </div>
  );
}
 
export default ColorChooserNode;
edges.ts
import { type Edge } from '@xyflow/react';
 
export const initialEdges = [
  { id: 'e1-2', source: '1', target: '2' },
  { id: 'e2-3', source: '2', target: '3' },
] as Edge[];
index.css
html,
body {
  margin: 0;
  font-family: sans-serif;
}
 
#app {
  width: 100vw;
  height: 100vh;
}
index.html
<!doctype html>
<html lang="en">
  <head>
    <meta charset="UTF-8" />
    <meta name="viewport" content="width=device-width, initial-scale=1.0" />
    <title>React Flow Example</title>
  </head>
  <body>
    <div id="app"></div>
    <script type="module" src="./index.tsx"></script>
  </body>
</html>
index.tsx
import { createRoot } from 'react-dom/client';
import App from './App';
 
import './index.css';
 
const container = document.querySelector('#app');
const root = createRoot(container);
 
root.render(<App />);
nodes.ts
import { type AppNode } from './types';
 
export const initialNodes = [
  {
    id: '1',
    type: 'colorChooser',
    data: { color: '#4FD1C5' },
    position: { x: 250, y: 25 },
  },
 
  {
    id: '2',
    type: 'colorChooser',
    data: { color: '#F6E05E' },
    position: { x: 100, y: 125 },
  },
  {
    id: '3',
    type: 'colorChooser',
    data: { color: '#B794F4' },
    position: { x: 250, y: 250 },
  },
] as AppNode[];
store.ts
import { create } from 'zustand';
import { addEdge, applyNodeChanges, applyEdgeChanges } from '@xyflow/react';
 
import { initialNodes } from './nodes';
import { initialEdges } from './edges';
import { type AppNode, type AppState, type ColorNode } from './types';
 
function isColorChooserNode(node: AppNode): node is ColorNode {
  return node.type === 'colorChooser';
}
 
// this is our useStore hook that we can use in our components to get parts of the store and call actions
const useStore = create<AppState>((set, get) => ({
  nodes: initialNodes,
  edges: initialEdges,
  onNodesChange: (changes) => {
    set({
      nodes: applyNodeChanges(changes, get().nodes),
    });
  },
  onEdgesChange: (changes) => {
    set({
      edges: applyEdgeChanges(changes, get().edges),
    });
  },
  onConnect: (connection) => {
    set({
      edges: addEdge(connection, get().edges),
    });
  },
  setNodes: (nodes) => {
    set({ nodes });
  },
  setEdges: (edges) => {
    set({ edges });
  },
  updateNodeColor: (nodeId, color) => {
    set({
      nodes: get().nodes.map((node) => {
        if (node.id === nodeId && isColorChooserNode(node)) {
          // it's important to create a new object here, to inform React Flow about the changes
          return { ...node, data: { ...node.data, color } };
        }
 
        return node;
      }),
    });
  },
}));
 
export default useStore;
types.ts
import {
  type Edge,
  type Node,
  type OnNodesChange,
  type OnEdgesChange,
  type OnConnect,
  type BuiltInNode,
} from '@xyflow/react';
 
export type ColorNode = Node<
  {
    color: string;
  },
  'colorChooser'
>;
 
export type AppNode = ColorNode | BuiltInNode;
 
export type AppState = {
  nodes: AppNode[];
  edges: Edge[];
  onNodesChange: OnNodesChange<AppNode>;
  onEdgesChange: OnEdgesChange;
  onConnect: OnConnect;
  setNodes: (nodes: AppNode[]) => void;
  setEdges: (edges: Edge[]) => void;
  updateNodeColor: (nodeId: string, color: string) => void;
};

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