Implementing Jotai for State Management

After installing Jotai, our first step will be to create the atoms for our state.

Creating Atoms for Cart and Reviews

Create a new atoms.ts file in the src/app/store directory. This file will contain the atoms for cart and reviews.

At the top of the file, import atom from Jotai, as well as the Cart and Review types from @/api/types:

import { atom } from "jotai";
import { Cart, Review } from "@/api/types";

First, we'll create the cart atom, which will have an initial state of an empty cart:

export const cartAtom = atom<Cart>({
	products: [],
});

The reviews atom can be either an array of Reviews or null, so we'll start it off at null:

export const reviewsAtom = atom<Review[] | null>(null);

With the atoms created, we'll move on to the store provider.

Creating a Store Provider

Create a new StoreProvider.tsx file inside of the store directory.

This will be a client component, and will use useState to track the store. We'll also bring in createStore as well as Provider from Jotai.

Jotai already comes with a provider, you just need to give it a store.

We'll initialize that store using useState, and then wrap our children in the Jotai Provider:

"use client";
import { useState } from "react";
import { createStore, Provider } from "jotai";

const StoreProvider = ({ children }: { children: React.ReactNode }) => {
	const [store] = useState(() => createStore());
	return <Provider store={store}>{children}</Provider>;
};

export default StoreProvider;

Now we'll wrap the layout with our new StoreProvider.

Wrap the Layout with the StoreProvider

Over in layout.tsx, import the StoreProvider and use it to replace the existing CartProvider from the context example.

We'll also give the Header and initial cart state, since the Header is where we'll initialize the store:

// inside the RootLayout return
<StoreProvider>
	<Header cart={cart} clearCartAction={clearCartAction} />
	<main className="mx-auto max-w-3xl">{children}</main>
</StoreProvider>

Initializing the Store in the Header

In Header.tsx, we'll use the useRef hook from React to track whether we've initialized our cart or not.

We'll also bring in useStore and useAtomValue hooks from Jotai. The useStore hook gets us access to the store, and useAtomValue gets us access to the value of an atom.

import { useState, useRef } from "react";
import { useStore, useAtomValue } from "jotai";

Inside the component, we'll create a store variable that comes from useStore(). We'll also create a loaded variable that uses useRef to check if we have initialized an atom or not.

If we haven't loaded yet, we will use store.set to set the cartAtom with the initial cart:

// inside the Header component:
const store = useStore();
const loaded = useRef(false);
if (!loaded.current) {
	store.set(cartAtom, initialCart);
	loaded.current = true;
}

Next we'll replace the existing cart variable from useCart with the one we'll get from calling useAtomValue with the cartAtom and the store to get it from:

const cart = useAtomValue(cartAtom, {
	store,
});

Now we can move on to updating the CartPopup component.

Updating the CartPopup Component

Inside of CartPopup.tsx, we'll import Jotai's useStore and useAtom hooks as well as our `cartAtom:

import { useStore, useAtom } from "jotai";
import { cartAtom } from "../store/atoms";

Like before, we'll replace the existing cart and setCart variables from useCart() with the ones we'll get from calling useAtom with the cartAtom and useStore():

const [cart, setCart] = useAtom(cartAtom, {
	store: useStore(),
});

Next we'll update the AddToCart component:

Updating the AddToCart Component

Inside of AddToCart.tsx, we'll import Jotai's useStore and useSetAtom hooks as well as our cartAtom.

import { useStore, useSetAtom } from "jotai";
import { cartAtom } from "../store/atoms";

We'll get setCart from useSetAtom with the cartAtom, and then set the store to the output of useStore():

const [cart, setCart] = useAtom(cartAtom, {
	store: useStore(),
});

With these changes in place, checking in the browser confirms that we have our initial state and are able to add items to the cart as expected.

Something to Watch For

Both the App Router and React documentation talk about never relying on the render order of components, but we're kind of doing that in our application.

Back in the layout.tsx file, remember that the Header and main are both siblings of each other, and the Header is doing the work of initializing the store.

In our case we are probably okay because the only thing that's over in the page is AddToCart, and AddToCart is only setting that store.

However, if this was a situation where this wouldn't be okay, one approach would be to create another client component which would take in children and do the initialization itself before doing any rendering.

That said, let's implement the review functionality.

Implementing Reviews

The process here will be similar to what we followed before.

Inside of AverageRating.tsx, we'll import useRef to check for initialization. We'll also bring in useStore and useAtomValue from Jotai, as well as our reviewsAtom:

// inside AverageRating.tsx
"use client";
import { useRef } from "react";
import { useStore, useAtomValue } from "jotai";

import { reviewsAtom } from "@/app/store/atoms";

Create a store variable that will be the result of calling useStore().

Then create an initialized variable by calling useRef with an initial value of false. If we aren't initialized, we'll call store.set() with the reviewsAtom and the initialReviews and toggle the initialized variable to true:

const store = useStore();
const initialized = useRef(false);
if (!initialized.current) {
	store.set(reviewsAtom, initialReviews);
	initialized.current = true;
}

Next we'll replace the reviews from useReviews with the useAtomValue that we'll pass the reviewsAtom into:

const reviews = useAtomValue(reviewsAtom, {
	store,
});

Note: If you get the atom value first, then you won't get the initial value! Instead, you'll get a copy of the original value of reviews, which is probably not what you want.

Now we can move on to the Reviews component.

Updating the Reviews Component

Over in Reviews.tsx we'll import useStore and useAtom from Jotai, as well as our reviewsAtom:

// inside Reviews.tsx
"use client";
import { useState, useRef } from "react";
import { useStore, useAtomValue } from "jotai";
import { Review } from "@/api/types";
import { reviewsAtom } from "@/app/store/atoms";

Now we can do our initialization just like we did before.

Bring in the store, check if it's initialized, and set the initialReviews if needed. Then replace the reviews and setReviews to come from a call to useAtom with the reviewsAtom and the store:

const store = useStore();
const initialized = useRef(false);
if (!initialized.current) {
	store.set(reviewsAtom, initialReviews);
	initialized.current = true;
}
const [reviews, setReviews] = useAtom(reviewsAtom, {
	store,
});

Recapping Our Work

Back in the browser, when we leave a review on a t-shirt it shows up and the average rating is updated accordingly. Like before, we can see that the reviews end up in the rendered source as well.

We've completed porting the app to use Jotai instead of React Context!

This process showed you how to manage atoms both at the layout and per-route levels, which you might consider using for your next App Router application.