Combining Jotai, LocalStorage, and React Query: A Powerful State Management Pattern

Table of Contents

I’ve been working with React applications for quite some time now, and state management has always been an interesting challenge. While there are many great solutions out there, I recently discovered a powerful pattern by combining three of my favorite tools: Jotai, LocalStorage, and React Query.

Why these three? Well, each solves a specific problem really well, but together they create something even more powerful.

Let me explain what brought me to this pattern. I was working on an application that needed to:

Load data from an API
Cache it locally for offline use
Keep the UI in sync with the latest data
Handle loading and error states gracefully

Sure, I could use Redux or Context API, but that felt like bringing a sledgehammer to crack a nut. That’s when I started playing with these tools.

The Problem #

While TanStack Query provides excellent data fetching and caching capabilities, and does support some persistence options, we’re looking for a more flexible atom-based solution that provides simpler state management. Here’s what we typically do with React Query:

function UserSettings() {
  const { data, isLoading } = useQuery({
    queryKey: ['settings'],
    queryFn: fetchUserSettings
  });

  // Basic caching but no flexible atom-based state management
  // Limited control over persistence
  // Less granular control over state updates
  
  return (
    <div>
      {isLoading ? <Loading /> : <SettingsForm data={data} />}
    </div>
  );
}

This works for data fetching, but when we need more granular control over our state and persistence, we can leverage Jotai’s atomic state management approach.

First Try: Just Jotai with Storage #

Jotai provides built-in storage utilities that make it easy to persist state:

import { atom } from 'jotai'
const settingsAtom = atomWithStorage('settings', null);

function UserSettings() {
  const [settings, setSettings] = useAtom(settingsAtom);
  
  useEffect(() => {
    fetchUserSettings().then(setSettings);
  }, []);

  // No automatic cache invalidation
  // No background updates
  // Manual error handling
  return <SettingsForm data={settings} />;
}

Better for local state, but we lose React Query’s powerful caching features. There must be a better way!

The Combined Solution #

Jotai’s Query extension provides a powerful integration with TanStack Query, giving us the best of both worlds:

TanStack Query’s powerful data fetching, caching, and synchronization
Jotai’s flexible atomic state management
Built-in storage persistence
Granular control over state updates
Optimized re-renders through atomic updates

Here’s how we can combine them:

import { atomWithStorage } from "jotai/utils";
import { atomWithQuery } from "jotai-tanstack-query";
import { atom } from "jotai";
import { QueryFunction, QueryKey } from "@tanstack/query-core";

type AtomWithQueryStorageOptions<T> = {
	key: string; // Unique key for local storage
	queryKey: QueryKey;
	queryFn: QueryFunction<T>;
	staleTimeMs?: number; // Duration after which data is considered stale (default: 1 hour)
};

export type AtomAction<T> =
	| { type: "refresh" }
	| { type: "invalidate" }
	| { type: "resetError" }
	| { type: "update"; newData: T };

/**
 * Creates an atom that combines Tanstack Query with local storage caching.
 *
 * @param options Configuration options
 * @param options.key Unique key for local storage
 * @param options.queryKey Tanstack Query key
 * @param options.queryFn Query function to fetch data
 * @param options.staleTimeMs Duration after which data is considered stale (default: 1 hour)
 *
 * @returns An atom that provides { data, error, loading } and supports actions:
 *          - refresh: Fetches fresh data
 *          - invalidate: Clears cached data
 *          - resetError: Clears error state
 *          - update: Manually updates data
 */
export const atomWithQueryStorage = <T>({
	key,
	queryKey,
	queryFn,
	staleTimeMs = 60 * 60 * 1000, // Default: 1 hour
}: AtomWithQueryStorageOptions<T>) => {
	// Atom for local storage persistence
	const storageAtom = atomWithStorage<{
		data: T | null | undefined;
		timestamp: number | null;
	}>(key, {
		data: null,
		timestamp: null,
	});

	// Atom for query error state
	const errorAtom = atom<Error | null>(null);

	// Atom for managing the query state
	const queryAtom = atomWithQuery((get) => ({
		queryKey,
		queryFn,
		staleTime: staleTimeMs, // Add this to control query staleness
		cacheTime: staleTimeMs * 2, // Optional: add cache time as well
		enabled:
			!get(storageAtom).data ||
			Date.now() - (get(storageAtom).timestamp || 0) > staleTimeMs,
	}));

	// Add a new atom for loading state
	const loadingAtom = atom<boolean>(false);

	// Combined atom for loading, error, and data states
	const combinedAtom = atom(
		(get) => {
			const { data, timestamp } = get(storageAtom);
			const error = get(errorAtom);
			const isLoading = get(loadingAtom);

			const isStale = !timestamp || Date.now() - timestamp > staleTimeMs;
			const queryResult = get(queryAtom);

			return {
				data: !isStale && data ? data : queryResult.data,
				error: error || queryResult.error, // Include both error sources
				loading: isLoading || (!data && isStale && !queryResult.data),
			};
		},
		(get, set, action: AtomAction<T>) => {
			if (action.type === "refresh") {
				set(loadingAtom, true); // Set loading before the query
				try {
					const freshData = get(queryAtom);
					set(storageAtom, { data: freshData.data, timestamp: Date.now() });
					set(errorAtom, freshData.error);
				} catch (error) {
					set(errorAtom, error as Error);
				} finally {
					set(loadingAtom, false); // Ensure loading is set to false after query
				}
			}

			if (action.type === "invalidate") {
				set(storageAtom, { data: null, timestamp: null }); // Clear storage
				set(loadingAtom, false);
				set(errorAtom, null);
			}

			if (action.type === "resetError") {
				set(errorAtom, null); // Reset error state
				set(loadingAtom, false);
			}

			if (action.type === "update") {
				set(storageAtom, { data: action.newData, timestamp: Date.now() }); // Manually update data
				set(loadingAtom, false);
				set(errorAtom, null);
			}
		},
	);

	return combinedAtom;
};

The beauty of this pattern is its simplicity. No complex setup, no boilerplate, just clean, functional code.

Now we can use it like this:

// Create the atom
const userSettingsAtom = atomWithQueryStorage({
  key: 'user-settings',
  queryKey: ['settings'],
  queryFn: fetchUserSettings,
  staleTimeMs: 24 * 60 * 60 * 1000 // 24 hours
});

// Use it in components
function UserSettings() {
  const [{ data, error, isLoading }, dispatch] = useAtom(userSettingsAtom);
  
  const handleRefresh = () => dispatch({ type: 'refresh' });
  const handleUpdate = (newData) => dispatch({ 
    type: 'update', 
    data: newData 
  });
  
  if (isLoading) return <Loading />;
  if (error) return <Error message={error.message} />;
  
  return (
    <div>
      <SettingsForm 
        data={data} 
        onUpdate={handleUpdate}
        onRefresh={handleRefresh}
      />
    </div>
  );
}

What I Learned #

This pattern taught me something important: sometimes the best solutions come from combining specialized tools:

TanStack Query excels at data fetching and caching
Jotai provides elegant atomic state management
Jotai Storage handles persistence
Jotai Query brings them all together

Together, they create a powerful, flexible system that’s still easy to understand and use.

Is this pattern perfect for every situation? Of course not. But for many common scenarios - especially when dealing with cached API data - it’s become my go-to solution.

Happy Coding!