ProNextJS
    solution

    Implementing Jotai for State Management

    Jack HerringtonJack Herrington

    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.

    Transcript

    Okay, so the first thing we're going to do is add the Jotai library to our example. All right, now the next thing we're going to do is remove the cart context, and we're going to go create another store directory. This one is going to have our atoms. So we'll bring in atom from Jotai.

    We'll bring in our types for cart and review, and then we'll create our cart atom that has this initial state of just an empty cart. And then the reviews atom that can be either an array of reviews or null, and we'll just start that off at null. Now to create our store that we will initialize on a per request basis over in the layout, we'll create a store provider.

    This is going to be a client component. We'll use useState to track the store. Then we'll bring in createStore as well as provider from Jotai. Jotai comes already with a provider, you just need to give it a store. So we're going to initialize that store using useState, and then wrap

    our children in the Jotai provided provider. Okay, let's go over to our layout and wrap our layout in our new provider. So now we'll go down here to cart provider, and we'll just replace that with store provider. Now who's going to initialize that store? Because that store needs to get initialized somewhere,

    so we're going to do that initialization in the header. So we'll add cart to our header, and then over in header we will say that we have an initial cart, and then we're going to bring in useRef. We're going to use useRef to track whether we have initialized our cart or not. Then we'll bring in useStore and useAtomValue

    from Jotai. useStore gets us access to the store, and useAtomValue gets us access to the value of an atom. There's also useAtom, which returns both the state as well as a setter, and useAtomSetter, I think, to get the setter for an atom if you just want that. All right, so now we're going to

    initialize our store. First we have to get access to the store, then we need to say are we loaded or not, and if we're not loaded then we want to initialize that atom, but we need access to the atom, so let's go bring that in. And then finally down here where we would use cart, we're going to use that atom value

    to get the cart atom value, and then as an option we'll give it the store, and that tells Jotai which store to use to go get that value of that atom. All right, with that out of the way, we can do it a little bit easier over in the cart pop-up. We'll bring in useStore and useAtom from Jotai.

    We'll bring in our cart atom from our atoms, and then to get cart and set cart we're going to call useAtom with that cart atom and with that store. So of course we got to get the store and we'll just call useStore. All right, now add to cart, let's bring

    in useStore and also useSetAtom. It wasn't useAtomSetter, useSetAtomMyBad. So we're going to get set cart from useSetAtom with our cart atom that we need to go get, and then we're going to set that store to the output of useStore.

    Seems to be doing okay. Let's go over and go to the dragon, add the dragon t-shirt to the cart, and away we go. How cool is that? Nice! All right, now one thing to watch out for if we look back in layout. So in here it's important to notice that the header and the main are both siblings of each other,

    and header is doing the work of actually initializing the store. Now in the AppWriter documentation and the React documentation they talk about never relying on the render order of components. We're kind of doing that here except for the fact that the only thing that's over in the page is add to cart, and add to carts

    only be setting that store. So I think in this case that's probably okay. If that wasn't okay, what you could do is probably create another client component that would take children. It itself would do the initialization first, and then any components within that would get rendered, and I think that would get around the problem of the ordering

    of the initialization of the atoms in the store. All right, so now let's do the reviews. We'll go into components, and we're going to do the thing where we initialize in each component. So we'll do initial reviews. We use the useRef to track whether we've been initialized or not, and then we'll get useStore and useAtomValue.

    Then we'll get our reviewsAtom, and then we'll get our initializer code. So we're going to bring in store with useStore. We're going to set the initialize to false so that it never runs again after that first render. Then we're going to set that reviewsAtom with our initial reviews, say that we're already done, and then

    after that we're going to get the atom value for the reviews. After that's really important, if you get the atom value first, then you won't get the initial value. You'll get a copy of the original value of reviews, which is probably not what you want. All right, let's finish up by implementing on reviews. Again, we'll bring in useRef, useStore, and useAtom,

    as well as the reviewsAtom. Again, we're going to call this initialReviews, and we're going to do our initialization. This is exactly the same as we had before, but of course we need call set reviews. So let's go do that down here. Okay, looks pretty good. Let's give it a try.

    Cool dragon shirt. One star. Not very good. All right, let's give it a try. Let's go and see that cool dragon shirt is in our SSR output. And there you go, a complete Jotai example of this port server-side rendering that shows you how to manage atoms

    both at the layout level as well as the per route level. A nice example to work off of, and an interesting state manager that you might want to take a look at for your next AppWriter application.