Use Jotai in Next.js app router

[ootkin]

Hi,

I would like to integrate Jotai in a next.js (v14) using app routers.

Following the Jotai documentation (even if the example reports page routers instead) I see that we need to use a global Provider.

I tried to implement the Provider in the RootLayout (/app/layout.tsx) but it seems like to work only when i refresh the page: when i change page using router.push() the state is not hydrated with useHydrateAtoms.

Do you have any ideas? Maybe the Provider must be placed in each page and in the root layout.

As the Next.js documentation says, the layout persist across routes and maintain state.

Thanks

[dai-shi]

Basically, the route change is not hydration. Hydration is only for the initial render.

So, we need to update with both useHydrateAtoms and useEffect for the route change.

...if I'm not mistaken.

The idea is this:

const useSyncAtom = (atom, value) => {
  useHydrateAtoms([atom, value]);
  const setAtom = useSetAtom(atom);
  useEffect(() => {
    setAtom(value);
  }, [setAtom, value]);
};

[ootkin]

Well, when i change navigation to a new Page (wich is a Server Components) i fetch some data and return them as a prop to the Page.

So a new state is computed server side and I think it is an hydration.

Right now I'm looking this

[ikaasraa]

I get this error

Error: Functions cannot be passed directly to Client Components unless you explicitly expose it by marking it with "use server". Or maybe you meant to call this function rather than return it.
  {toString: ..., init: "", read: ..., write: function defaultWrite}

I use nextjs 14 and approuter

[ootkin]

Furthermore, if I use useHydrateAtoms within the Jotai Provider, I've got this error:

Detected multiple Jotai instances. It may cause unexpected behavior with the default store. #2044

[dai-shi]

If you see the error without HMR, it sounds bad. You are using Provider right? You may want to use Provider separately for different routes.

[ootkin]

I use HMR and I use Provider separately for different routes: even after I removed useHydrateAtoms.

[MiroslavPetrik]

@ootkin it's a feature of the useHydrateAtoms to allow only one hydration.

Let's say you have a page like app/[userId]/page.tsx with a jotai provider in the app/layout.tsx.
When you start by navigating to /1, the hydration will work.
Next, if you router.push("/2") the hydration won't work second time for the same atom.

This might be confusing, as you are correct - the SSR takes place, so we expect hydration... but since the router.push is client side navigation, we must keep in mind that the app has a local state in the memory.. Looking into the useHydrateAtoms source, we can see that there is a local hydratedSet created for each store..

This means, that if you have <Provider> in your RootLayout, there is a shared store for all of your routes. And so the atom hydration works per app, not per route because the hydratedSet stays the same...

So the solution is to use the <Provider> in your page.tsx.

Other (hypothetical) solution could be to call hydratedSet.delete for all the atoms you wish to rehydrated, just before the router.push.
That would require changes to the util, as the hydratedSet is currently not accessible... perhaps the hook could return evictAtoms cleanup function.

const {evict} = useHydrateAtoms([[atom, value]]);

This could help for other client side navigation cases like router.refresh() which is usually called after a form submit.
Here, we need to reinitialize/rehydrate the form, which won't happen despite the page having a local <Provider>.
So invalidating the cache (the hydratedSet here) of already hydrated atoms could help.

evict() // enables subsequent hydration
router.refresh() // fetches new data & merges the RSC
// form will be reinitialized 

With such cache eviction we won't need to set the atoms manually after the hydration

useEffect(() => {
   setAtom(value);
 }, [setAtom, value]);

This to me looks like a "second hydration" and can be confusing, or easy to forget.

@dai-shi thoughts?

[dai-shi]

useEffect in a for-loop can violate the rules of hooks. The length of atomValues can vary.

The trick in useHydrateAtoms is that it's only effective for uninitialized atoms.

[Tomaszal]

Right, for some reason I thought it wouldn't be a problem since useSyncAtoms is a hook itself, but yeah of course the array can still change.. Still I think there's a way to do it without an extra atom and memoization like so (unless I'm missing something obvious again):

const useSyncAtoms = (atomValues) => {
  useHydrateAtoms(atomValues);
  const store = useStore();
  useEffect(
    () => atomValues.forEach(([atom, value]) => store.set(atom, value)),
    [store, ...atomValues.flat()],
  );
};

[dai-shi]

This is not a Jotai issue, but a React issue. You can't safely do useEffect(fn, [...arr]) or just useEffect(fn, arr) unless arr is stable. If the length of arr changes, React will warn you in dev mode.

[Tomaszal]

I realize this isn't a Jotai issue, wasn't implying that. I'm just looking for a way to do this in a wrapper that doesn't require memoizing each atom's value before passing it to useRehydrateAtoms. I guess the simplest way would be to use a wrapper for individual atoms (as in #2337 (comment)) instead of specifying an array of tuples, I was just trying to keep the API similar to the existing useHydrateAtoms, but seems like there isn't a clean way to do that.

[dai-shi]

Dealing with arrays in React is always a bit tricky.

[dai-shi]

related: #669

[benjamin-andersen]

#2337 (comment)

@ootkin I believe you can use the root template for the app router in Nextjs. It should re-create the <Provider> anytime a user navigates between routes. This will allow the hydration to occur every time the route changes, and has the added benefit of not needing to wrap every page in the <Provider>.
https://nextjs.org/docs/app/building-your-application/routing/layouts-and-templates#templates

I did some rudimentary testing with this and it seemed to fix the issue for me.

// app/template.tsx

import {Provider} from 'jotai'

interface RootTemplateProps {
    children: React.ReactNode
}

const RootTemplate: React.FC<RootTemplateProps> = ({children}) => (
    <Provider>
        {children}
    </Provider>
)

export default RootTemplate

If you want to synchronize the hydration with cache invalidation (e.g. revalidateTag()), then you will need to use something like this hook instead of useHydrateAtoms(), since the hydration occurs only once per route change.