6 files changed, 68 insertions, 134 deletions

diff --git a/src/client/atoms/cards.ts b/src/client/atoms/cards.ts
index f053ab9..c0275b0 100644
--- a/src/client/atoms/cards.ts
+++ b/src/client/atoms/cards.ts
@@ -1,5 +1,6 @@
+import { atomFamily } from "jotai/utils";
+import { atomWithSuspenseQuery } from "jotai-tanstack-query";
 import { apiClient } from "../api/client";
-import { createReloadableAtomFamily } from "./utils";
 
 export interface Card {
 	id: string;
@@ -20,12 +21,15 @@ export interface Card {
 // Cards by Deck - Suspense-compatible
 // =====================
 
-export const cardsByDeckAtomFamily = createReloadableAtomFamily(
-	async (deckId: string) => {
-		const res = await apiClient.rpc.api.decks[":deckId"].cards.$get({
-			param: { deckId },
-		});
-		const data = await apiClient.handleResponse<{ cards: Card[] }>(res);
-		return data.cards;
-	},
+export const cardsByDeckAtomFamily = atomFamily((deckId: string) =>
+	atomWithSuspenseQuery(() => ({
+		queryKey: ["decks", deckId, "cards"],
+		queryFn: async () => {
+			const res = await apiClient.rpc.api.decks[":deckId"].cards.$get({
+				param: { deckId },
+			});
+			const data = await apiClient.handleResponse<{ cards: Card[] }>(res);
+			return data.cards;
+		},
+	})),
 );
diff --git a/src/client/atoms/decks.ts b/src/client/atoms/decks.ts
index e9b0d03..fc68727 100644
--- a/src/client/atoms/decks.ts
+++ b/src/client/atoms/decks.ts
@@ -1,5 +1,6 @@
+import { atomFamily } from "jotai/utils";
+import { atomWithSuspenseQuery } from "jotai-tanstack-query";
 import { apiClient } from "../api/client";
-import { createReloadableAtom, createReloadableAtomFamily } from "./utils";
 
 export interface Deck {
 	id: string;
@@ -15,24 +16,30 @@ export interface Deck {
 // Decks List - Suspense-compatible
 // =====================
 
-export const decksAtom = createReloadableAtom(async () => {
-	const res = await apiClient.rpc.api.decks.$get(undefined, {
-		headers: apiClient.getAuthHeader(),
-	});
-	const data = await apiClient.handleResponse<{ decks: Deck[] }>(res);
-	return data.decks;
-});
+export const decksAtom = atomWithSuspenseQuery(() => ({
+	queryKey: ["decks"],
+	queryFn: async () => {
+		const res = await apiClient.rpc.api.decks.$get(undefined, {
+			headers: apiClient.getAuthHeader(),
+		});
+		const data = await apiClient.handleResponse<{ decks: Deck[] }>(res);
+		return data.decks;
+	},
+}));
 
 // =====================
 // Single Deck by ID - Suspense-compatible
 // =====================
 
-export const deckByIdAtomFamily = createReloadableAtomFamily(
-	async (deckId: string) => {
-		const res = await apiClient.rpc.api.decks[":id"].$get({
-			param: { id: deckId },
-		});
-		const data = await apiClient.handleResponse<{ deck: Deck }>(res);
-		return data.deck;
-	},
+export const deckByIdAtomFamily = atomFamily((deckId: string) =>
+	atomWithSuspenseQuery(() => ({
+		queryKey: ["decks", deckId],
+		queryFn: async () => {
+			const res = await apiClient.rpc.api.decks[":id"].$get({
+				param: { id: deckId },
+			});
+			const data = await apiClient.handleResponse<{ deck: Deck }>(res);
+			return data.deck;
+		},
+	})),
 );
diff --git a/src/client/atoms/index.ts b/src/client/atoms/index.ts
index 1e13222..1de2544 100644
--- a/src/client/atoms/index.ts
+++ b/src/client/atoms/index.ts
@@ -37,6 +37,3 @@ export {
 	syncStatusAtom,
 	useSyncInit,
 } from "./sync";
-
-// Utilities
-export { createReloadableAtom, createReloadableAtomFamily } from "./utils";
diff --git a/src/client/atoms/noteTypes.ts b/src/client/atoms/noteTypes.ts
index adc9d44..fb99b14 100644
--- a/src/client/atoms/noteTypes.ts
+++ b/src/client/atoms/noteTypes.ts
@@ -1,5 +1,5 @@
+import { atomWithSuspenseQuery } from "jotai-tanstack-query";
 import { apiClient } from "../api/client";
-import { createReloadableAtom } from "./utils";
 
 export interface NoteType {
 	id: string;
@@ -15,8 +15,11 @@ export interface NoteType {
 // NoteTypes List - Suspense-compatible
 // =====================
 
-export const noteTypesAtom = createReloadableAtom(async () => {
-	const res = await apiClient.rpc.api["note-types"].$get();
-	const data = await apiClient.handleResponse<{ noteTypes: NoteType[] }>(res);
-	return data.noteTypes;
-});
+export const noteTypesAtom = atomWithSuspenseQuery(() => ({
+	queryKey: ["noteTypes"],
+	queryFn: async () => {
+		const res = await apiClient.rpc.api["note-types"].$get();
+		const data = await apiClient.handleResponse<{ noteTypes: NoteType[] }>(res);
+		return data.noteTypes;
+	},
+}));
diff --git a/src/client/atoms/study.ts b/src/client/atoms/study.ts
index 2e3e1ea..1911737 100644
--- a/src/client/atoms/study.ts
+++ b/src/client/atoms/study.ts
@@ -1,6 +1,7 @@
+import { atomFamily } from "jotai/utils";
+import { atomWithSuspenseQuery } from "jotai-tanstack-query";
 import { apiClient } from "../api/client";
 import { shuffle } from "../utils/shuffle";
-import { createReloadableAtomFamily } from "./utils";
 
 export interface StudyCard {
 	id: string;
@@ -36,24 +37,27 @@ export interface StudyData {
 // Study Session - Suspense-compatible
 // =====================
 
-export const studyDataAtomFamily = createReloadableAtomFamily(
-	async (deckId: string): Promise<StudyData> => {
-		// Fetch deck and due cards in parallel
-		const [deckRes, cardsRes] = await Promise.all([
-			apiClient.rpc.api.decks[":id"].$get({ param: { id: deckId } }),
-			apiClient.rpc.api.decks[":deckId"].study.$get({ param: { deckId } }),
-		]);
+export const studyDataAtomFamily = atomFamily((deckId: string) =>
+	atomWithSuspenseQuery(() => ({
+		queryKey: ["decks", deckId, "study"],
+		queryFn: async (): Promise<StudyData> => {
+			// Fetch deck and due cards in parallel
+			const [deckRes, cardsRes] = await Promise.all([
+				apiClient.rpc.api.decks[":id"].$get({ param: { id: deckId } }),
+				apiClient.rpc.api.decks[":deckId"].study.$get({ param: { deckId } }),
+			]);
 
-		const deckData = await apiClient.handleResponse<{ deck: StudyDeck }>(
-			deckRes,
-		);
-		const cardsData = await apiClient.handleResponse<{ cards: StudyCard[] }>(
-			cardsRes,
-		);
+			const deckData = await apiClient.handleResponse<{ deck: StudyDeck }>(
+				deckRes,
+			);
+			const cardsData = await apiClient.handleResponse<{
+				cards: StudyCard[];
+			}>(cardsRes);
 
-		return {
-			deck: deckData.deck,
-			cards: shuffle(cardsData.cards),
-		};
-	},
+			return {
+				deck: deckData.deck,
+				cards: shuffle(cardsData.cards),
+			};
+		},
+	})),
 );
diff --git a/src/client/atoms/utils.ts b/src/client/atoms/utils.ts
deleted file mode 100644
index e7af288..0000000
--- a/src/client/atoms/utils.ts
+++ /dev/null
@@ -1,81 +0,0 @@
-import { atom, type Getter, type WritableAtom } from "jotai";
-
-// Symbol to identify reload action
-const RELOAD = Symbol("reload");
-
-/**
- * A WritableAtom that returns T (or Promise<T> before hydration) and accepts
- * an optional T value for hydration, or undefined to trigger reload.
- */
-export type ReloadableAtom<T> = WritableAtom<T | Promise<T>, [T?], void>;
-
-/**
- * Creates an async atom that can be reloaded by calling its setter.
- * Read the atom to get the data (suspends while loading).
- * Set the atom with no args to trigger a reload.
- * Set the atom with a value to hydrate (useful for testing).
- */
-export function createReloadableAtom<T>(
-	getter: (get: Getter) => Promise<T>,
-): ReloadableAtom<T> {
-	const refetchKeyAtom = atom(0);
-	// Stores hydrated value - undefined means not hydrated
-	const hydratedValueAtom = atom<{ value: T } | undefined>(undefined);
-
-	return atom(
-		// Not using async here - returns T synchronously when hydrated, Promise<T> when fetching
-		(get): T | Promise<T> => {
-			// Check for hydrated value first (sync path - avoids Suspense)
-			const hydrated = get(hydratedValueAtom);
-			if (hydrated !== undefined) {
-				return hydrated.value;
-			}
-			// Async path - will trigger Suspense
-			get(refetchKeyAtom);
-			return getter(get);
-		},
-		(_get, set, action?: T | typeof RELOAD) => {
-			if (action === undefined || action === RELOAD) {
-				// Trigger reload: clear hydrated value and bump refetch key
-				set(hydratedValueAtom, undefined);
-				set(refetchKeyAtom, (k) => k + 1);
-			} else {
-				// Hydrate with value
-				set(hydratedValueAtom, { value: action });
-			}
-		},
-	);
-}
-
-// Track all atom family caches for test cleanup
-const atomFamilyCaches: Map<unknown, unknown>[] = [];
-
-/**
- * Creates a reloadable atom family for parameterized async data.
- * Each unique parameter gets its own cached atom with reload capability.
- */
-export function createReloadableAtomFamily<T, P extends string | number>(
-	getter: (param: P, get: Getter) => Promise<T>,
-): (param: P) => ReloadableAtom<T> {
-	const cache = new Map<P, ReloadableAtom<T>>();
-	atomFamilyCaches.push(cache);
-
-	return (param: P): ReloadableAtom<T> => {
-		let reloadableAtom = cache.get(param);
-		if (!reloadableAtom) {
-			reloadableAtom = createReloadableAtom((get) => getter(param, get));
-			cache.set(param, reloadableAtom);
-		}
-		return reloadableAtom;
-	};
-}
-
-/**
- * Clears all atom family caches. Call this in test beforeEach/afterEach
- * to ensure tests don't share cached atoms.
- */
-export function clearAtomFamilyCaches() {
-	for (const cache of atomFamilyCaches) {
-		cache.clear();
-	}
-}