Creating Auth management with vite, zustand, tanstack router and tanstack query. Is the implementation correct?

Question

Alright so the goal I had was to create an auth layer that would manage the state and protected routes globally for my vite project. I am using jwt tokens and I want to create a smooth and optimized setup.

I am using zustand, tanstack query and tanstack router in my Vite + React (typescript) project.

---

Let me start with the AuthStore. This code sets up the state variables I want to use along with the roles, errors etc. I have 3 functions: login, logout, and session for now.

export const useAuthStore = createWithEqualityFn<AuthStore>((set) => ({

    user: null, isAuthenticated: null, isHost: false, loading: false, error: null, info: "",

    login: async (userLogIn) => {
        set({ loading: true, error: null });
        try {
            const response = await auth.login(userLogIn);
            set({ user: response.user, isAuthenticated: true, isHost: response.isHost });
        } catch (error) {
            const err = handleError(error);
            if (err.statusCode === 403) {
                set({ error: { message: "User is banned", statusCode: 403 } });
            } else {
                set({ error: err });
            }
        } finally { set({ loading: false }) }
    },

    logout: () => {
        set({ loading: true, error: null });
        try{
            auth.logout();
        } catch (error) {
            set({ error: handleError(error) });
        } finally { set({ isAuthenticated: false, isHost: false, loading: false, user: null }) }
    },

    session: async () => {
        set({ loading: true, error: null });
        try {
            const response = await auth.session();
            set({ user: response.user, isAuthenticated: true, isHost: response.isHost });
        } catch (error) {
            set({ error: handleError(error) });
            throw error;
        } finally { set({ loading: false }) }
    },
}));

Note that isHost is a role as I have to do role based routing as well. Also I will fix error handling later. handleError is just a wrapper that does the same. auth.session() is a get function that is similar to a pseudo login with tokens.

I wanted to use the session function to determine if the user is logged in or logged out when the page loads. The jwt token management is handled with cookies. Since isAuthenticated is in a zustand store, it can be used anywhere in the SPA.

---

This session function is used in getAuth function.

export async function checkAuth() {
  const { isAuthenticated, session } = useAuthStore.getState();

  if (isAuthenticated === null) {
    try {
        await session();
    } catch (error) {
      console.log("Session failed. Please login again.");
      useAuthStore.setState({ isAuthenticated: false });
    }
  }
}

With tanstack router, I am using route-tree file based routing. So in the __root.tsx file,

export const Route = createRootRoute({
    beforeLoad: checkAuth,
    component: () => (
        <>
            <Navbar />
            <Outlet />
            <TanStackRouterDevtools />
        </>
    ),
})

So when the page first loads, the isAuthenticated is determined as it goes from null to true or false and I can manage all the protected routes as follows:

export const Route = createFileRoute('/profile/')({
    beforeLoad: () => {
        const isAuthenticated = useAuthStore.getState().isAuthenticated;
        if (!isAuthenticated) {
            throw redirect({ to: '/login' });
        }
    },
    component: RouteComponent,
})

This would make /profile a protected route that needs isAuthenticated to be true else it would redirect to /login.

---

This would also mean that for UI components such as navbar, I can use the zustand store values. Here is an example

export default function Navbar() {
    const isAuthenticated = useAuthStore((state) => state.isAuthenticated);
    return (
        <nav>
            <button>{isAuthenticated? Profile : Login}</button>
        </nav>
    )
}

---

With this setup I am able to achieve what I set out to do - global auth state management. The reason I am posting this is because I couldn't find any similar implementations and I do not know if this is bug free, device-compatible or logically correct. Any suggestions would be greatly appreciated.