From 6dd30ee85ecf3bfc08e854a54987305967522777 Mon Sep 17 00:00:00 2001
From: Christian Baroni <7061887+christianbaroni@users.noreply.github.com>
Date: Wed, 11 Dec 2024 21:05:52 +0000
Subject: [PATCH 01/29] Add createRemoteRainbowStore

---
 src/state/internal/createRainbowStore.ts      |   5 +-
 .../internal/createRemoteRainbowStore.ts      | 323 ++++++++++++++++++
 .../internal/tests/RainbowRemoteStoreTest.tsx | 113 ++++++
 3 files changed, 438 insertions(+), 3 deletions(-)
 create mode 100644 src/state/internal/createRemoteRainbowStore.ts
 create mode 100644 src/state/internal/tests/RainbowRemoteStoreTest.tsx

diff --git a/src/state/internal/createRainbowStore.ts b/src/state/internal/createRainbowStore.ts
index 9b7d9a39dd7..57a760f69d1 100644
--- a/src/state/internal/createRainbowStore.ts
+++ b/src/state/internal/createRainbowStore.ts
@@ -1,8 +1,7 @@
 import { debounce } from 'lodash';
 import { MMKV } from 'react-native-mmkv';
-import { create } from 'zustand';
+import { StateCreator, create } from 'zustand';
 import { PersistOptions, StorageValue, persist, subscribeWithSelector } from 'zustand/middleware';
-import { StateCreator } from 'zustand/vanilla';
 import { RainbowError, logger } from '@/logger';
 
 const PERSIST_RATE_LIMIT_MS = 3000;
@@ -12,7 +11,7 @@ const rainbowStorage = new MMKV({ id: 'rainbow-storage' });
 /**
  * Configuration options for creating a persistable Rainbow store.
  */
-interface RainbowPersistConfig<S> {
+export interface RainbowPersistConfig<S> {
   /**
    * A function to convert the serialized string back into the state object.
    * If not provided, the default deserializer is used.
diff --git a/src/state/internal/createRemoteRainbowStore.ts b/src/state/internal/createRemoteRainbowStore.ts
new file mode 100644
index 00000000000..6fc47e3e9c8
--- /dev/null
+++ b/src/state/internal/createRemoteRainbowStore.ts
@@ -0,0 +1,323 @@
+/* eslint-disable @typescript-eslint/no-explicit-any */
+
+import { StateCreator, StoreApi, UseBoundStore, create } from 'zustand';
+import { subscribeWithSelector } from 'zustand/middleware';
+import { createRainbowStore, RainbowPersistConfig } from './createRainbowStore';
+
+/**
+ * Represents the status of the remote data fetching process.
+ */
+type RemoteStatus = 'idle' | 'loading' | 'success' | 'error';
+
+/**
+ * Configuration options for remote data fetching.
+ */
+interface FetchOptions {
+  cacheTime?: number;
+  force?: boolean;
+  staleTime?: number;
+}
+
+/**
+ * Represents a cached query result.
+ */
+interface CacheEntry<TData> {
+  data: TData;
+  lastFetchedAt: number;
+}
+
+/**
+ * The base store state including remote fields and actions.
+ */
+type StoreState<TData, TParams extends Record<string, any>> = {
+  data: TData | null;
+  enabled: boolean;
+  error: Error | null;
+  lastFetchedAt: number | null;
+  queryCache: Record<string, CacheEntry<TData>>;
+  status: RemoteStatus;
+  subscriptionCount: number;
+  fetch: (params?: TParams, options?: FetchOptions) => Promise<void>;
+  isDataExpired: (cacheTimeOverride?: number) => boolean;
+  isStale: (staleTimeOverride?: number) => boolean;
+  reset: () => void;
+};
+
+/**
+ * A specialized store interface combining Zustand's store API with remote fetching.
+ */
+export interface RemoteStore<TData, TParams extends Record<string, any>, S extends StoreState<TData, TParams>>
+  extends UseBoundStore<StoreApi<S>> {
+  enabled: boolean;
+  fetch: (params?: TParams, options?: FetchOptions) => Promise<void>;
+  isDataExpired: (override?: number) => boolean;
+  isStale: (override?: number) => boolean;
+  reset: () => void;
+}
+
+/**
+ * Configuration options for creating a remote-enabled Rainbow store.
+ */
+type RemoteRainbowStoreConfig<TQueryFnData, TParams extends Record<string, any>, TData, S extends StoreState<TData, TParams>> = {
+  cacheTime?: number;
+  defaultParams?: TParams;
+  disableDataCache?: boolean;
+  enabled?: boolean;
+  queryKey: readonly unknown[] | ((params: TParams) => readonly unknown[]);
+  staleTime?: number;
+  fetcher: (params: TParams) => TQueryFnData | Promise<TQueryFnData>;
+  onFetched?: (data: TData, store: RemoteStore<TData, TParams, S>) => void;
+  transform?: (data: TQueryFnData) => TData;
+};
+
+const SEVEN_DAYS = 1000 * 60 * 60 * 24 * 7;
+const TWO_MINUTES = 1000 * 60 * 2;
+
+/**
+ * Creates a remote-enabled Rainbow store with data fetching capabilities.
+ *
+ * We use a `U` generic to represent user-defined additional state, and then define:
+ * S = StoreState<TData, TParams> & U
+ *
+ * This ensures that the base fields are always present, and user-added fields are merged in.
+ */
+export function createRemoteRainbowStore<
+  TQueryFnData,
+  TParams extends Record<string, any> = Record<string, any>,
+  U = unknown,
+  TData = TQueryFnData,
+>(
+  config: RemoteRainbowStoreConfig<TQueryFnData, TParams, TData, StoreState<TData, TParams> & U>,
+  customStateCreator?: StateCreator<U, [], [['zustand/subscribeWithSelector', never]]>,
+  persistConfig?: RainbowPersistConfig<StoreState<TData, TParams> & U>
+): RemoteStore<TData, TParams, StoreState<TData, TParams> & U> {
+  type S = StoreState<TData, TParams> & U;
+
+  const {
+    cacheTime = SEVEN_DAYS,
+    defaultParams,
+    disableDataCache = true,
+    enabled = true,
+    queryKey,
+    staleTime = TWO_MINUTES,
+    fetcher,
+    onFetched,
+    transform,
+  } = config;
+
+  let activeFetchPromise: Promise<void> | null = null;
+  let activeRefetchTimeout: NodeJS.Timeout | null = null;
+  let lastFetchKey: string | null = null;
+
+  const getQueryKey = (params: TParams): string => {
+    const key = typeof queryKey === 'function' ? queryKey(params) : queryKey;
+    return JSON.stringify(key);
+  };
+
+  const initialData: Omit<S, 'fetch' | 'isStale' | 'isDataExpired' | 'reset'> = {
+    data: null,
+    enabled,
+    error: null,
+    lastFetchedAt: null,
+    queryCache: {},
+    status: 'idle',
+    subscriptionCount: 0,
+  } as unknown as Omit<S, 'fetch' | 'isStale' | 'isDataExpired' | 'reset'>;
+
+  const createState: StateCreator<S, [], [['zustand/subscribeWithSelector', never]]> = (set, get, api) => {
+    let isRefetchScheduled = false;
+
+    const pruneCache = (state: S): S => {
+      if (disableDataCache) return state;
+      const now = Date.now();
+      const newCache: Record<string, CacheEntry<TData>> = {};
+      Object.entries(state.queryCache).forEach(([key, entry]) => {
+        if (now - entry.lastFetchedAt <= cacheTime) {
+          newCache[key] = entry;
+        }
+      });
+      return { ...state, queryCache: newCache };
+    };
+
+    const scheduleNextFetch = (params: TParams) => {
+      if (isRefetchScheduled || staleTime <= 0) return;
+      if (activeRefetchTimeout) clearTimeout(activeRefetchTimeout);
+
+      isRefetchScheduled = true;
+      activeRefetchTimeout = setTimeout(() => {
+        isRefetchScheduled = false;
+        const store = get();
+        if (store.subscriptionCount > 0) {
+          store.fetch(params, { force: true });
+        }
+      }, staleTime);
+    };
+
+    const baseMethods = {
+      async fetch(params: TParams | undefined, options: FetchOptions | undefined) {
+        if (!get().enabled) return;
+
+        const effectiveParams = params ?? defaultParams ?? ({} as TParams);
+        const currentQueryKey = getQueryKey(effectiveParams);
+
+        if (activeFetchPromise && lastFetchKey === currentQueryKey && get().status === 'loading' && !options?.force) {
+          return activeFetchPromise;
+        }
+
+        if (!options?.force && !disableDataCache) {
+          const currentState = get();
+          const cached = currentState.queryCache[currentQueryKey];
+          if (cached && Date.now() - cached.lastFetchedAt <= (options?.staleTime ?? staleTime)) {
+            set(state => ({ ...state, data: cached.data }));
+            return;
+          }
+        }
+
+        set(state => ({ ...state, status: 'loading', error: null }));
+        lastFetchKey = currentQueryKey;
+
+        const fetchOperation = async () => {
+          try {
+            const rawData = await fetcher(effectiveParams);
+            const transformedData = transform ? transform(rawData) : (rawData as TData);
+
+            set(state => {
+              const newState = {
+                ...state,
+                error: null,
+                lastFetchedAt: Date.now(),
+                status: 'success' as const,
+              };
+
+              if (!disableDataCache) {
+                newState.queryCache = {
+                  ...newState.queryCache,
+                  [currentQueryKey]: {
+                    data: transformedData,
+                    lastFetchedAt: Date.now(),
+                  },
+                };
+              }
+
+              if (!onFetched) newState.data = transformedData;
+
+              return pruneCache(newState);
+            });
+
+            scheduleNextFetch(effectiveParams);
+
+            if (onFetched) {
+              onFetched(transformedData, remoteCapableStore);
+            }
+          } catch (error: any) {
+            console.log('[ERROR]:', error);
+            set(state => ({ ...state, error, status: 'error' as const }));
+            scheduleNextFetch(effectiveParams);
+          } finally {
+            activeFetchPromise = null;
+            lastFetchKey = null;
+          }
+        };
+
+        activeFetchPromise = fetchOperation();
+        return activeFetchPromise;
+      },
+
+      isStale(staleTimeOverride?: number) {
+        const { lastFetchedAt } = get();
+        const effectiveStaleTime = staleTimeOverride ?? staleTime;
+        if (lastFetchedAt === null) return true;
+        return Date.now() - lastFetchedAt > effectiveStaleTime;
+      },
+
+      isDataExpired(cacheTimeOverride?: number) {
+        const { lastFetchedAt } = get();
+        const effectiveCacheTime = cacheTimeOverride ?? cacheTime;
+        if (lastFetchedAt === null) return true;
+        return Date.now() - lastFetchedAt > effectiveCacheTime;
+      },
+
+      reset() {
+        if (activeRefetchTimeout) {
+          clearTimeout(activeRefetchTimeout);
+          activeRefetchTimeout = null;
+        }
+        activeFetchPromise = null;
+        lastFetchKey = null;
+        isRefetchScheduled = false;
+        set(initialData as Partial<S> as S);
+      },
+    };
+
+    // If customStateCreator is provided, it will return user-defined fields (U)
+    const userState = customStateCreator?.(set, get, api) ?? ({} as U);
+
+    const subscribeWithSelector = api.subscribe;
+    api.subscribe = (listener: (state: S, prevState: S) => void) => {
+      set(prev => ({ ...prev, subscriptionCount: prev.subscriptionCount + 1 }));
+      const unsubscribe = subscribeWithSelector(listener);
+
+      const handleSetEnabled = subscribeWithSelector((state: S, prev: S) => {
+        if (state.enabled !== prev.enabled) {
+          if (state.enabled) {
+            if (!state.data || state.isStale()) {
+              state.fetch(defaultParams);
+            } else {
+              scheduleNextFetch(defaultParams ?? ({} as TParams));
+            }
+          } else {
+            if (activeRefetchTimeout) {
+              clearTimeout(activeRefetchTimeout);
+              activeRefetchTimeout = null;
+            }
+            isRefetchScheduled = false;
+          }
+        }
+      });
+
+      const currentState = get();
+      if (!currentState.data || currentState.isStale()) {
+        currentState.fetch(defaultParams, { force: true });
+      } else {
+        scheduleNextFetch(defaultParams ?? ({} as TParams));
+      }
+
+      return () => {
+        handleSetEnabled();
+        unsubscribe();
+        set(prev => {
+          const newCount = Math.max(prev.subscriptionCount - 1, 0);
+          if (newCount === 0) {
+            if (activeRefetchTimeout) {
+              clearTimeout(activeRefetchTimeout);
+              activeRefetchTimeout = null;
+            }
+            isRefetchScheduled = false;
+          }
+          return { ...prev, subscriptionCount: newCount };
+        });
+      };
+    };
+
+    // Merge base data, user state, and methods into the final store state
+    return {
+      ...initialData,
+      ...userState,
+      ...baseMethods,
+    } satisfies S;
+  };
+
+  const baseStore = persistConfig?.storageKey
+    ? createRainbowStore<StoreState<TData, TParams> & U>(createState, persistConfig)
+    : create<StoreState<TData, TParams> & U>()(subscribeWithSelector(createState));
+
+  const remoteCapableStore: RemoteStore<TData, TParams, S> = Object.assign(baseStore, {
+    enabled,
+    fetch: (params?: TParams, options?: FetchOptions) => baseStore.getState().fetch(params, options),
+    isDataExpired: (override?: number) => baseStore.getState().isDataExpired(override),
+    isStale: (override?: number) => baseStore.getState().isStale(override),
+    reset: () => baseStore.getState().reset(),
+  });
+
+  return remoteCapableStore;
+}
diff --git a/src/state/internal/tests/RainbowRemoteStoreTest.tsx b/src/state/internal/tests/RainbowRemoteStoreTest.tsx
new file mode 100644
index 00000000000..42d5b32a215
--- /dev/null
+++ b/src/state/internal/tests/RainbowRemoteStoreTest.tsx
@@ -0,0 +1,113 @@
+import React, { memo, useEffect } from 'react';
+import { StyleSheet, View } from 'react-native';
+import { Address } from 'viem';
+import ButtonPressAnimation from '@/components/animations/ButtonPressAnimation';
+import { Text } from '@/design-system';
+import { SupportedCurrencyKey } from '@/references';
+import { queryUserAssets } from '@/__swaps__/screens/Swap/resources/assets/userAssets';
+import { ParsedAssetsDictByChain } from '@/__swaps__/types/assets';
+import { createRemoteRainbowStore } from '../internal/createRemoteRainbowStore';
+
+function getRandomAddress() {
+  return Math.random() < 0.5 ? '0x2e67869829c734ac13723A138a952F7A8B56e774' : '0xCFB83E14AEd465c79F3F82f4cfF4ff7965897644';
+}
+
+type QueryParams = { address: Address; currency: SupportedCurrencyKey };
+
+type UserAssetsState = {
+  userAssets: ParsedAssetsDictByChain;
+  getHighestValueAsset: () => number;
+};
+
+export const userAssetsStore = createRemoteRainbowStore<ParsedAssetsDictByChain, QueryParams, UserAssetsState, ParsedAssetsDictByChain>(
+  {
+    enabled: true,
+    queryKey: ['userAssets'],
+    // staleTime: 5000, // 5s
+    staleTime: 30 * 60 * 1000, // 30m
+
+    fetcher: (/* { address, currency } */) => queryUserAssets({ address: getRandomAddress(), currency: 'USD' }),
+    onFetched: (data, store) => store.setState({ data }),
+    transform: data => {
+      const lastFetchedAt = Date.now();
+      const formattedTimeWithSeconds = lastFetchedAt
+        ? new Date(lastFetchedAt).toLocaleTimeString('en-US', { hour: '2-digit', minute: '2-digit', second: '2-digit' })
+        : 'N/A';
+      console.log('[Transform - Last Fetch Attempt]: ', formattedTimeWithSeconds);
+
+      return data;
+    },
+  },
+  (set, get) => ({
+    userAssets: [],
+
+    getHighestValueAsset: () => {
+      const data = get().userAssets;
+      const highestValueAsset = Object.values(data)
+        .flatMap(chainAssets => Object.values(chainAssets))
+        .reduce((max, asset) => {
+          return Math.max(max, Number(asset.balance.display));
+        }, 0);
+      return highestValueAsset;
+    },
+
+    setUserAssets: (data: ParsedAssetsDictByChain) => set({ userAssets: data }),
+  }),
+  {
+    storageKey: 'userAssetsTesting79876',
+  }
+);
+
+export const UserAssetsTest = memo(function UserAssetsTest() {
+  const data = userAssetsStore(state => state.data);
+  const enabled = userAssetsStore(state => state.enabled);
+
+  console.log('RERENDER');
+
+  useEffect(() => {
+    const status = userAssetsStore.getState().status;
+    const isFetching = status === 'loading';
+    // eslint-disable-next-line no-nested-ternary
+    const emojiForStatus = isFetching ? '🔄' : status === 'success' ? '✅' : '❌';
+    console.log('[NEW STATUS]:', emojiForStatus, status);
+
+    if (data) {
+      const allTokens = Object.values(data).flatMap(chainAssets => Object.values(chainAssets));
+      const first5Tokens = allTokens.slice(0, 5);
+      console.log('[First 5 Token Symbols]:', first5Tokens.map(token => token.symbol).join(', '));
+    }
+  }, [data]);
+
+  return (
+    data && (
+      <View style={styles.container}>
+        <Text color="label" size="17pt" weight="heavy">
+          Number of assets: {Object.values(data).reduce((acc, chainAssets) => acc + Object.keys(chainAssets).length, 0)}
+        </Text>
+        <ButtonPressAnimation onPress={() => userAssetsStore.setState({ enabled: !enabled })} style={styles.button}>
+          <Text color="label" size="17pt" weight="heavy">
+            {enabled ? 'Disable fetching' : 'Enable fetching'}
+          </Text>
+        </ButtonPressAnimation>
+      </View>
+    )
+  );
+});
+
+const styles = StyleSheet.create({
+  button: {
+    alignItems: 'center',
+    backgroundColor: 'blue',
+    borderRadius: 22,
+    height: 44,
+    justifyContent: 'center',
+    paddingHorizontal: 20,
+  },
+  container: {
+    alignItems: 'center',
+    flex: 1,
+    flexDirection: 'column',
+    gap: 32,
+    justifyContent: 'center',
+  },
+});

From 68101ed6dbd57da7e8f985ffe913b76165ec53a9 Mon Sep 17 00:00:00 2001
From: Christian Baroni <7061887+christianbaroni@users.noreply.github.com>
Date: Thu, 12 Dec 2024 04:45:49 +0000
Subject: [PATCH 02/29] Fix import path

---
 src/state/internal/tests/RainbowRemoteStoreTest.tsx | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/state/internal/tests/RainbowRemoteStoreTest.tsx b/src/state/internal/tests/RainbowRemoteStoreTest.tsx
index 42d5b32a215..8321d496d79 100644
--- a/src/state/internal/tests/RainbowRemoteStoreTest.tsx
+++ b/src/state/internal/tests/RainbowRemoteStoreTest.tsx
@@ -6,7 +6,7 @@ import { Text } from '@/design-system';
 import { SupportedCurrencyKey } from '@/references';
 import { queryUserAssets } from '@/__swaps__/screens/Swap/resources/assets/userAssets';
 import { ParsedAssetsDictByChain } from '@/__swaps__/types/assets';
-import { createRemoteRainbowStore } from '../internal/createRemoteRainbowStore';
+import { createRemoteRainbowStore } from '../createRemoteRainbowStore';
 
 function getRandomAddress() {
   return Math.random() < 0.5 ? '0x2e67869829c734ac13723A138a952F7A8B56e774' : '0xCFB83E14AEd465c79F3F82f4cfF4ff7965897644';

From 8a4882729450322abe58d7bc0db94487001afa2b Mon Sep 17 00:00:00 2001
From: Christian Baroni <7061887+christianbaroni@users.noreply.github.com>
Date: Thu, 12 Dec 2024 21:38:29 +0000
Subject: [PATCH 03/29] Improve types, rename to createRainbowQueryStore

---
 ...bowStore.ts => createRainbowQueryStore.ts} | 89 +++++++++++--------
 ...toreTest.tsx => RainbowQueryStoreTest.tsx} | 41 +++++----
 2 files changed, 74 insertions(+), 56 deletions(-)
 rename src/state/internal/{createRemoteRainbowStore.ts => createRainbowQueryStore.ts} (79%)
 rename src/state/internal/tests/{RainbowRemoteStoreTest.tsx => RainbowQueryStoreTest.tsx} (75%)

diff --git a/src/state/internal/createRemoteRainbowStore.ts b/src/state/internal/createRainbowQueryStore.ts
similarity index 79%
rename from src/state/internal/createRemoteRainbowStore.ts
rename to src/state/internal/createRainbowQueryStore.ts
index 6fc47e3e9c8..91743e0216a 100644
--- a/src/state/internal/createRemoteRainbowStore.ts
+++ b/src/state/internal/createRainbowQueryStore.ts
@@ -2,6 +2,8 @@
 
 import { StateCreator, StoreApi, UseBoundStore, create } from 'zustand';
 import { subscribeWithSelector } from 'zustand/middleware';
+import { IS_DEV } from '@/env';
+import { logger, RainbowError } from '@/logger';
 import { createRainbowStore, RainbowPersistConfig } from './createRainbowStore';
 
 /**
@@ -27,7 +29,7 @@ interface CacheEntry<TData> {
 }
 
 /**
- * The base store state including remote fields and actions.
+ * The base store state including query-related fields and actions.
  */
 type StoreState<TData, TParams extends Record<string, any>> = {
   data: TData | null;
@@ -46,7 +48,7 @@ type StoreState<TData, TParams extends Record<string, any>> = {
 /**
  * A specialized store interface combining Zustand's store API with remote fetching.
  */
-export interface RemoteStore<TData, TParams extends Record<string, any>, S extends StoreState<TData, TParams>>
+export interface QueryStore<TData, TParams extends Record<string, any>, S extends StoreState<TData, TParams>>
   extends UseBoundStore<StoreApi<S>> {
   enabled: boolean;
   fetch: (params?: TParams, options?: FetchOptions) => Promise<void>;
@@ -58,53 +60,60 @@ export interface RemoteStore<TData, TParams extends Record<string, any>, S exten
 /**
  * Configuration options for creating a remote-enabled Rainbow store.
  */
-type RemoteRainbowStoreConfig<TQueryFnData, TParams extends Record<string, any>, TData, S extends StoreState<TData, TParams>> = {
+type RainbowQueryStoreConfig<TQueryFnData, TParams extends Record<string, any>, TData, S extends StoreState<TData, TParams>> = {
+  fetcher: (params: TParams) => TQueryFnData | Promise<TQueryFnData>;
+  onFetched?: (data: TData, store: QueryStore<TData, TParams, S>) => void;
+  transform?: (data: TQueryFnData) => TData;
   cacheTime?: number;
   defaultParams?: TParams;
   disableDataCache?: boolean;
   enabled?: boolean;
   queryKey: readonly unknown[] | ((params: TParams) => readonly unknown[]);
   staleTime?: number;
-  fetcher: (params: TParams) => TQueryFnData | Promise<TQueryFnData>;
-  onFetched?: (data: TData, store: RemoteStore<TData, TParams, S>) => void;
-  transform?: (data: TQueryFnData) => TData;
 };
 
 const SEVEN_DAYS = 1000 * 60 * 60 * 24 * 7;
 const TWO_MINUTES = 1000 * 60 * 2;
+const FIVE_SECONDS = 1000 * 5;
+const MIN_STALE_TIME = FIVE_SECONDS;
 
 /**
  * Creates a remote-enabled Rainbow store with data fetching capabilities.
- *
- * We use a `U` generic to represent user-defined additional state, and then define:
- * S = StoreState<TData, TParams> & U
- *
- * This ensures that the base fields are always present, and user-added fields are merged in.
+ * @template TQueryFnData - The raw data type returned by the fetcher
+ * @template TParams - Parameters passed to the fetcher function
+ * @template TData - The transformed data type (defaults to TQueryFnData)
+ * @template U - Additional user-defined state
  */
-export function createRemoteRainbowStore<
+export function createRainbowQueryStore<
   TQueryFnData,
   TParams extends Record<string, any> = Record<string, any>,
   U = unknown,
   TData = TQueryFnData,
 >(
-  config: RemoteRainbowStoreConfig<TQueryFnData, TParams, TData, StoreState<TData, TParams> & U>,
+  config: RainbowQueryStoreConfig<TQueryFnData, TParams, TData, StoreState<TData, TParams> & U>,
   customStateCreator?: StateCreator<U, [], [['zustand/subscribeWithSelector', never]]>,
   persistConfig?: RainbowPersistConfig<StoreState<TData, TParams> & U>
-): RemoteStore<TData, TParams, StoreState<TData, TParams> & U> {
+): QueryStore<TData, TParams, StoreState<TData, TParams> & U> {
   type S = StoreState<TData, TParams> & U;
 
   const {
+    fetcher,
+    onFetched,
+    transform,
     cacheTime = SEVEN_DAYS,
     defaultParams,
     disableDataCache = true,
     enabled = true,
     queryKey,
     staleTime = TWO_MINUTES,
-    fetcher,
-    onFetched,
-    transform,
   } = config;
 
+  if (IS_DEV && staleTime < MIN_STALE_TIME) {
+    console.warn(
+      `[RainbowQueryStore${persistConfig?.storageKey ? `: ${persistConfig.storageKey}` : ''}] ❌ Stale times under ${MIN_STALE_TIME / 1000} seconds are not recommended.`
+    );
+  }
+
   let activeFetchPromise: Promise<void> | null = null;
   let activeRefetchTimeout: NodeJS.Timeout | null = null;
   let lastFetchKey: string | null = null;
@@ -114,15 +123,15 @@ export function createRemoteRainbowStore<
     return JSON.stringify(key);
   };
 
-  const initialData: Omit<S, 'fetch' | 'isStale' | 'isDataExpired' | 'reset'> = {
+  const initialData = {
     data: null,
     enabled,
     error: null,
     lastFetchedAt: null,
     queryCache: {},
-    status: 'idle',
+    status: 'idle' as const,
     subscriptionCount: 0,
-  } as unknown as Omit<S, 'fetch' | 'isStale' | 'isDataExpired' | 'reset'>;
+  };
 
   const createState: StateCreator<S, [], [['zustand/subscribeWithSelector', never]]> = (set, get, api) => {
     let isRefetchScheduled = false;
@@ -146,9 +155,8 @@ export function createRemoteRainbowStore<
       isRefetchScheduled = true;
       activeRefetchTimeout = setTimeout(() => {
         isRefetchScheduled = false;
-        const store = get();
-        if (store.subscriptionCount > 0) {
-          store.fetch(params, { force: true });
+        if (get().subscriptionCount > 0) {
+          get().fetch(params, { force: true });
         }
       }, staleTime);
     };
@@ -165,20 +173,20 @@ export function createRemoteRainbowStore<
         }
 
         if (!options?.force && !disableDataCache) {
-          const currentState = get();
-          const cached = currentState.queryCache[currentQueryKey];
+          const cached = get().queryCache[currentQueryKey];
           if (cached && Date.now() - cached.lastFetchedAt <= (options?.staleTime ?? staleTime)) {
             set(state => ({ ...state, data: cached.data }));
             return;
           }
         }
 
-        set(state => ({ ...state, status: 'loading', error: null }));
+        set(state => ({ ...state, error: null, status: 'loading' }));
         lastFetchKey = currentQueryKey;
 
         const fetchOperation = async () => {
           try {
-            const rawData = await fetcher(effectiveParams);
+            const result = await fetcher(effectiveParams);
+            const rawData = result instanceof Promise ? await result : result;
             const transformedData = transform ? transform(rawData) : (rawData as TData);
 
             set(state => {
@@ -207,10 +215,14 @@ export function createRemoteRainbowStore<
             scheduleNextFetch(effectiveParams);
 
             if (onFetched) {
-              onFetched(transformedData, remoteCapableStore);
+              onFetched(transformedData, queryCapableStore);
             }
-          } catch (error: any) {
-            console.log('[ERROR]:', error);
+          } catch (error) {
+            logger.error(
+              new RainbowError(`[createRainbowQueryStore: ${persistConfig?.storageKey || currentQueryKey}]: Failed to fetch data`),
+              { error }
+            );
+            // TODO: Improve retry logic
             set(state => ({ ...state, error, status: 'error' as const }));
             scheduleNextFetch(effectiveParams);
           } finally {
@@ -245,7 +257,7 @@ export function createRemoteRainbowStore<
         activeFetchPromise = null;
         lastFetchKey = null;
         isRefetchScheduled = false;
-        set(initialData as Partial<S> as S);
+        set(state => ({ ...state, ...initialData }));
       },
     };
 
@@ -275,9 +287,10 @@ export function createRemoteRainbowStore<
         }
       });
 
-      const currentState = get();
-      if (!currentState.data || currentState.isStale()) {
-        currentState.fetch(defaultParams, { force: true });
+      const { data, fetch, isStale } = get();
+
+      if (!data || isStale()) {
+        fetch(defaultParams, { force: true });
       } else {
         scheduleNextFetch(defaultParams ?? ({} as TParams));
       }
@@ -304,20 +317,20 @@ export function createRemoteRainbowStore<
       ...initialData,
       ...userState,
       ...baseMethods,
-    } satisfies S;
+    };
   };
 
   const baseStore = persistConfig?.storageKey
     ? createRainbowStore<StoreState<TData, TParams> & U>(createState, persistConfig)
     : create<StoreState<TData, TParams> & U>()(subscribeWithSelector(createState));
 
-  const remoteCapableStore: RemoteStore<TData, TParams, S> = Object.assign(baseStore, {
-    enabled,
+  const queryCapableStore: QueryStore<TData, TParams, S> = Object.assign(baseStore, {
     fetch: (params?: TParams, options?: FetchOptions) => baseStore.getState().fetch(params, options),
     isDataExpired: (override?: number) => baseStore.getState().isDataExpired(override),
     isStale: (override?: number) => baseStore.getState().isStale(override),
     reset: () => baseStore.getState().reset(),
+    enabled,
   });
 
-  return remoteCapableStore;
+  return queryCapableStore;
 }
diff --git a/src/state/internal/tests/RainbowRemoteStoreTest.tsx b/src/state/internal/tests/RainbowQueryStoreTest.tsx
similarity index 75%
rename from src/state/internal/tests/RainbowRemoteStoreTest.tsx
rename to src/state/internal/tests/RainbowQueryStoreTest.tsx
index 8321d496d79..a3e22886ab2 100644
--- a/src/state/internal/tests/RainbowRemoteStoreTest.tsx
+++ b/src/state/internal/tests/RainbowQueryStoreTest.tsx
@@ -6,7 +6,7 @@ import { Text } from '@/design-system';
 import { SupportedCurrencyKey } from '@/references';
 import { queryUserAssets } from '@/__swaps__/screens/Swap/resources/assets/userAssets';
 import { ParsedAssetsDictByChain } from '@/__swaps__/types/assets';
-import { createRemoteRainbowStore } from '../createRemoteRainbowStore';
+import { createRainbowQueryStore } from '../createRainbowQueryStore';
 
 function getRandomAddress() {
   return Math.random() < 0.5 ? '0x2e67869829c734ac13723A138a952F7A8B56e774' : '0xCFB83E14AEd465c79F3F82f4cfF4ff7965897644';
@@ -14,29 +14,29 @@ function getRandomAddress() {
 
 type QueryParams = { address: Address; currency: SupportedCurrencyKey };
 
-type UserAssetsState = {
+type TestStore = {
   userAssets: ParsedAssetsDictByChain;
   getHighestValueAsset: () => number;
+  setUserAssets: (data: ParsedAssetsDictByChain) => void;
 };
 
-export const userAssetsStore = createRemoteRainbowStore<ParsedAssetsDictByChain, QueryParams, UserAssetsState, ParsedAssetsDictByChain>(
+export const userAssetsStore = createRainbowQueryStore<ParsedAssetsDictByChain, QueryParams, TestStore, ParsedAssetsDictByChain>(
   {
-    enabled: true,
-    queryKey: ['userAssets'],
-    // staleTime: 5000, // 5s
-    staleTime: 30 * 60 * 1000, // 30m
-
-    fetcher: (/* { address, currency } */) => queryUserAssets({ address: getRandomAddress(), currency: 'USD' }),
-    onFetched: (data, store) => store.setState({ data }),
+    fetcher: () => queryUserAssets({ address: getRandomAddress(), currency: 'USD' }),
+    // onFetched: (data, store) => store.setState({ userAssets: data }),
     transform: data => {
-      const lastFetchedAt = Date.now();
-      const formattedTimeWithSeconds = lastFetchedAt
-        ? new Date(lastFetchedAt).toLocaleTimeString('en-US', { hour: '2-digit', minute: '2-digit', second: '2-digit' })
-        : 'N/A';
+      const formattedTimeWithSeconds = new Date(Date.now()).toLocaleTimeString('en-US', {
+        hour: '2-digit',
+        minute: '2-digit',
+        second: '2-digit',
+      });
       console.log('[Transform - Last Fetch Attempt]: ', formattedTimeWithSeconds);
-
       return data;
     },
+
+    disableDataCache: false,
+    queryKey: ['userAssets'],
+    staleTime: 30 * 60 * 1000, // 30m
   },
   (set, get) => ({
     userAssets: [],
@@ -54,6 +54,7 @@ export const userAssetsStore = createRemoteRainbowStore<ParsedAssetsDictByChain,
     setUserAssets: (data: ParsedAssetsDictByChain) => set({ userAssets: data }),
   }),
   {
+    // partialize: state => ({ userAssets: state.userAssets }),
     storageKey: 'userAssetsTesting79876',
   }
 );
@@ -62,7 +63,7 @@ export const UserAssetsTest = memo(function UserAssetsTest() {
   const data = userAssetsStore(state => state.data);
   const enabled = userAssetsStore(state => state.enabled);
 
-  console.log('RERENDER');
+  console.log('RERENDER - enabled:', enabled);
 
   useEffect(() => {
     const status = userAssetsStore.getState().status;
@@ -72,8 +73,9 @@ export const UserAssetsTest = memo(function UserAssetsTest() {
     console.log('[NEW STATUS]:', emojiForStatus, status);
 
     if (data) {
-      const allTokens = Object.values(data).flatMap(chainAssets => Object.values(chainAssets));
-      const first5Tokens = allTokens.slice(0, 5);
+      const first5Tokens = Object.values(data)
+        .flatMap(chainAssets => Object.values(chainAssets))
+        .slice(0, 5);
       console.log('[First 5 Token Symbols]:', first5Tokens.map(token => token.symbol).join(', '));
     }
   }, [data]);
@@ -94,6 +96,9 @@ export const UserAssetsTest = memo(function UserAssetsTest() {
   );
 });
 
+const initialData = userAssetsStore.getState().data;
+console.log('[Initial Data Exists]:', !!initialData);
+
 const styles = StyleSheet.create({
   button: {
     alignItems: 'center',

From c8121cb936bc7578b40469d737fa41159dc766d7 Mon Sep 17 00:00:00 2001
From: Christian Baroni <7061887+christianbaroni@users.noreply.github.com>
Date: Thu, 12 Dec 2024 21:48:12 +0000
Subject: [PATCH 04/29] Clean up createRainbowStore

---
 src/state/internal/createRainbowStore.ts | 12 ++++++------
 1 file changed, 6 insertions(+), 6 deletions(-)

diff --git a/src/state/internal/createRainbowStore.ts b/src/state/internal/createRainbowStore.ts
index 57a760f69d1..371095b6c00 100644
--- a/src/state/internal/createRainbowStore.ts
+++ b/src/state/internal/createRainbowStore.ts
@@ -17,6 +17,11 @@ export interface RainbowPersistConfig<S> {
    * If not provided, the default deserializer is used.
    */
   deserializer?: (serializedState: string) => StorageValue<Partial<S>>;
+  /**
+   * A function to perform persisted state migration.
+   * This function will be called when persisted state versions mismatch with the one specified here.
+   */
+  migrate?: (persistedState: unknown, version: number) => S | Promise<S>;
   /**
    * A function that determines which parts of the state should be persisted.
    * By default, the entire state is persisted.
@@ -37,11 +42,6 @@ export interface RainbowPersistConfig<S> {
    * @default 0
    */
   version?: number;
-  /**
-   * A function to perform persisted state migration.
-   * This function will be called when persisted state versions mismatch with the one specified here.
-   */
-  migrate?: (persistedState: unknown, version: number) => S | Promise<S>;
 }
 
 /**
@@ -151,11 +151,11 @@ export function createRainbowStore<S = unknown>(
   return create<S>()(
     subscribeWithSelector(
       persist(createState, {
+        migrate: persistConfig.migrate,
         name: persistConfig.storageKey,
         partialize: persistConfig.partialize || (state => state),
         storage: persistStorage,
         version,
-        migrate: persistConfig.migrate,
       })
     )
   );

From 3fa324fa6090d2e4470603950b3a2c724795bf5a Mon Sep 17 00:00:00 2001
From: Christian Baroni <7061887+christianbaroni@users.noreply.github.com>
Date: Fri, 13 Dec 2024 05:01:42 +0000
Subject: [PATCH 05/29] Ensure internal state is persisted when partialize is
 used, remove any types

---
 src/state/internal/createRainbowQueryStore.ts | 61 ++++++++++++++++---
 1 file changed, 54 insertions(+), 7 deletions(-)

diff --git a/src/state/internal/createRainbowQueryStore.ts b/src/state/internal/createRainbowQueryStore.ts
index 91743e0216a..34e53e8fd18 100644
--- a/src/state/internal/createRainbowQueryStore.ts
+++ b/src/state/internal/createRainbowQueryStore.ts
@@ -1,5 +1,3 @@
-/* eslint-disable @typescript-eslint/no-explicit-any */
-
 import { StateCreator, StoreApi, UseBoundStore, create } from 'zustand';
 import { subscribeWithSelector } from 'zustand/middleware';
 import { IS_DEV } from '@/env';
@@ -31,7 +29,7 @@ interface CacheEntry<TData> {
 /**
  * The base store state including query-related fields and actions.
  */
-type StoreState<TData, TParams extends Record<string, any>> = {
+type StoreState<TData, TParams extends Record<string, unknown>> = {
   data: TData | null;
   enabled: boolean;
   error: Error | null;
@@ -45,10 +43,15 @@ type StoreState<TData, TParams extends Record<string, any>> = {
   reset: () => void;
 };
 
+/**
+ * The keys that make up the internal state of the store.
+ */
+type InternalStateKey = keyof StoreState<unknown, Record<string, unknown>>;
+
 /**
  * A specialized store interface combining Zustand's store API with remote fetching.
  */
-export interface QueryStore<TData, TParams extends Record<string, any>, S extends StoreState<TData, TParams>>
+export interface QueryStore<TData, TParams extends Record<string, unknown>, S extends StoreState<TData, TParams>>
   extends UseBoundStore<StoreApi<S>> {
   enabled: boolean;
   fetch: (params?: TParams, options?: FetchOptions) => Promise<void>;
@@ -60,7 +63,7 @@ export interface QueryStore<TData, TParams extends Record<string, any>, S extend
 /**
  * Configuration options for creating a remote-enabled Rainbow store.
  */
-type RainbowQueryStoreConfig<TQueryFnData, TParams extends Record<string, any>, TData, S extends StoreState<TData, TParams>> = {
+type RainbowQueryStoreConfig<TQueryFnData, TParams extends Record<string, unknown>, TData, S extends StoreState<TData, TParams>> = {
   fetcher: (params: TParams) => TQueryFnData | Promise<TQueryFnData>;
   onFetched?: (data: TData, store: QueryStore<TData, TParams, S>) => void;
   transform?: (data: TQueryFnData) => TData;
@@ -77,6 +80,8 @@ const TWO_MINUTES = 1000 * 60 * 2;
 const FIVE_SECONDS = 1000 * 5;
 const MIN_STALE_TIME = FIVE_SECONDS;
 
+const DISCARDABLE_INTERNAL_STATE: InternalStateKey[] = ['fetch', 'isDataExpired', 'isStale', 'reset', 'subscriptionCount'];
+
 /**
  * Creates a remote-enabled Rainbow store with data fetching capabilities.
  * @template TQueryFnData - The raw data type returned by the fetcher
@@ -86,7 +91,7 @@ const MIN_STALE_TIME = FIVE_SECONDS;
  */
 export function createRainbowQueryStore<
   TQueryFnData,
-  TParams extends Record<string, any> = Record<string, any>,
+  TParams extends Record<string, unknown> = Record<string, unknown>,
   U = unknown,
   TData = TQueryFnData,
 >(
@@ -320,8 +325,15 @@ export function createRainbowQueryStore<
     };
   };
 
+  const combinedPersistConfig = persistConfig
+    ? {
+        ...persistConfig,
+        partialize: createBlendedPartialize(persistConfig.partialize),
+      }
+    : undefined;
+
   const baseStore = persistConfig?.storageKey
-    ? createRainbowStore<StoreState<TData, TParams> & U>(createState, persistConfig)
+    ? createRainbowStore<StoreState<TData, TParams> & U>(createState, combinedPersistConfig)
     : create<StoreState<TData, TParams> & U>()(subscribeWithSelector(createState));
 
   const queryCapableStore: QueryStore<TData, TParams, S> = Object.assign(baseStore, {
@@ -334,3 +346,38 @@ export function createRainbowQueryStore<
 
   return queryCapableStore;
 }
+
+/**
+ * Checks whether a state key is internal and should be discarded from persistence.
+ */
+function shouldDiscardInternalState(key: InternalStateKey | string): key is InternalStateKey {
+  return DISCARDABLE_INTERNAL_STATE.includes(key as InternalStateKey);
+}
+
+/**
+ * Creates a combined partialize function that ensures internal query state is always
+ * persisted while respecting user-defined persistence preferences.
+ */
+function createBlendedPartialize<TData, TParams extends Record<string, unknown>, S extends StoreState<TData, TParams> & U, U = unknown>(
+  userPartialize: ((state: StoreState<TData, TParams> & U) => Partial<StoreState<TData, TParams> & U>) | undefined
+) {
+  return (state: S) => {
+    const internalStateToPersist = {
+      data: state.data,
+      enabled: state.enabled,
+      error: state.error,
+      lastFetchedAt: state.lastFetchedAt,
+      queryCache: state.queryCache,
+      status: state.status,
+    };
+
+    for (const key in state) {
+      if (shouldDiscardInternalState(key)) delete state[key];
+    }
+
+    return {
+      ...(userPartialize ? userPartialize(state) : state),
+      ...internalStateToPersist,
+    };
+  };
+}

From 938c7ded3729115845115be0ca2b2dd09fba5453 Mon Sep 17 00:00:00 2001
From: Christian Baroni <7061887+christianbaroni@users.noreply.github.com>
Date: Fri, 13 Dec 2024 06:15:04 +0000
Subject: [PATCH 06/29] Make createBlendedPartialize faster, improve typing

---
 src/state/internal/createRainbowQueryStore.ts | 45 +++++++++++--------
 1 file changed, 26 insertions(+), 19 deletions(-)

diff --git a/src/state/internal/createRainbowQueryStore.ts b/src/state/internal/createRainbowQueryStore.ts
index 34e53e8fd18..772d408ebef 100644
--- a/src/state/internal/createRainbowQueryStore.ts
+++ b/src/state/internal/createRainbowQueryStore.ts
@@ -80,7 +80,25 @@ const TWO_MINUTES = 1000 * 60 * 2;
 const FIVE_SECONDS = 1000 * 5;
 const MIN_STALE_TIME = FIVE_SECONDS;
 
-const DISCARDABLE_INTERNAL_STATE: InternalStateKey[] = ['fetch', 'isDataExpired', 'isStale', 'reset', 'subscriptionCount'];
+/**
+ * A map of internal state keys to whether they should be included in the persisted state.
+ */
+const SHOULD_PERSIST_INTERNAL_STATE_MAP: Record<string, boolean> = {
+  /* State to persist */
+  data: true,
+  enabled: true,
+  error: true,
+  lastFetchedAt: true,
+  queryCache: true,
+  status: true,
+
+  /* State and methods to discard */
+  fetch: false,
+  isDataExpired: false,
+  isStale: false,
+  reset: false,
+  subscriptionCount: false,
+} satisfies Record<InternalStateKey, boolean>;
 
 /**
  * Creates a remote-enabled Rainbow store with data fetching capabilities.
@@ -266,7 +284,7 @@ export function createRainbowQueryStore<
       },
     };
 
-    // If customStateCreator is provided, it will return user-defined fields (U)
+    /* If customStateCreator is provided, it will return user-defined fields (U) */
     const userState = customStateCreator?.(set, get, api) ?? ({} as U);
 
     const subscribeWithSelector = api.subscribe;
@@ -347,13 +365,6 @@ export function createRainbowQueryStore<
   return queryCapableStore;
 }
 
-/**
- * Checks whether a state key is internal and should be discarded from persistence.
- */
-function shouldDiscardInternalState(key: InternalStateKey | string): key is InternalStateKey {
-  return DISCARDABLE_INTERNAL_STATE.includes(key as InternalStateKey);
-}
-
 /**
  * Creates a combined partialize function that ensures internal query state is always
  * persisted while respecting user-defined persistence preferences.
@@ -362,22 +373,18 @@ function createBlendedPartialize<TData, TParams extends Record<string, unknown>,
   userPartialize: ((state: StoreState<TData, TParams> & U) => Partial<StoreState<TData, TParams> & U>) | undefined
 ) {
   return (state: S) => {
-    const internalStateToPersist = {
-      data: state.data,
-      enabled: state.enabled,
-      error: state.error,
-      lastFetchedAt: state.lastFetchedAt,
-      queryCache: state.queryCache,
-      status: state.status,
-    };
+    const internalStateToPersist: Partial<S> = {};
 
     for (const key in state) {
-      if (shouldDiscardInternalState(key)) delete state[key];
+      if (key in SHOULD_PERSIST_INTERNAL_STATE_MAP) {
+        if (SHOULD_PERSIST_INTERNAL_STATE_MAP[key]) internalStateToPersist[key] = state[key];
+        delete state[key];
+      }
     }
 
     return {
       ...(userPartialize ? userPartialize(state) : state),
       ...internalStateToPersist,
-    };
+    } satisfies Partial<S>;
   };
 }

From 4a6fc530e6786f4064e8e2a4ba8193d9dcb542b6 Mon Sep 17 00:00:00 2001
From: Christian Baroni <7061887+christianbaroni@users.noreply.github.com>
Date: Fri, 13 Dec 2024 08:17:35 +0000
Subject: [PATCH 07/29] Improve persist map readability

---
 src/state/internal/createRainbowQueryStore.ts | 33 +++++++++----------
 1 file changed, 16 insertions(+), 17 deletions(-)

diff --git a/src/state/internal/createRainbowQueryStore.ts b/src/state/internal/createRainbowQueryStore.ts
index 772d408ebef..395a2766c7c 100644
--- a/src/state/internal/createRainbowQueryStore.ts
+++ b/src/state/internal/createRainbowQueryStore.ts
@@ -80,24 +80,23 @@ const TWO_MINUTES = 1000 * 60 * 2;
 const FIVE_SECONDS = 1000 * 5;
 const MIN_STALE_TIME = FIVE_SECONDS;
 
-/**
- * A map of internal state keys to whether they should be included in the persisted state.
- */
+const [persist, discard] = [true, false];
+
 const SHOULD_PERSIST_INTERNAL_STATE_MAP: Record<string, boolean> = {
-  /* State to persist */
-  data: true,
-  enabled: true,
-  error: true,
-  lastFetchedAt: true,
-  queryCache: true,
-  status: true,
-
-  /* State and methods to discard */
-  fetch: false,
-  isDataExpired: false,
-  isStale: false,
-  reset: false,
-  subscriptionCount: false,
+  /* Internal state to persist if the store is persisted */
+  data: persist,
+  enabled: persist,
+  error: persist,
+  lastFetchedAt: persist,
+  queryCache: persist,
+  status: persist,
+
+  /* Internal state and methods to discard */
+  fetch: discard,
+  isDataExpired: discard,
+  isStale: discard,
+  reset: discard,
+  subscriptionCount: discard,
 } satisfies Record<InternalStateKey, boolean>;
 
 /**

From d2c00610c62bad27350a271ead09bafef978d65d Mon Sep 17 00:00:00 2001
From: Christian Baroni <7061887+christianbaroni@users.noreply.github.com>
Date: Mon, 16 Dec 2024 19:42:08 +0000
Subject: [PATCH 08/29] Add dynamic params, API cleanup

---
 src/state/internal/createRainbowQueryStore.ts | 275 ++++++++++++------
 src/state/internal/createStore.ts             |   3 +
 src/state/internal/signal.ts                  | 165 +++++++++++
 .../internal/tests/RainbowQueryStoreTest.tsx  | 156 ++++++----
 4 files changed, 466 insertions(+), 133 deletions(-)
 create mode 100644 src/state/internal/signal.ts

diff --git a/src/state/internal/createRainbowQueryStore.ts b/src/state/internal/createRainbowQueryStore.ts
index 395a2766c7c..bb39bd07df7 100644
--- a/src/state/internal/createRainbowQueryStore.ts
+++ b/src/state/internal/createRainbowQueryStore.ts
@@ -3,11 +3,21 @@ import { subscribeWithSelector } from 'zustand/middleware';
 import { IS_DEV } from '@/env';
 import { logger, RainbowError } from '@/logger';
 import { createRainbowStore, RainbowPersistConfig } from './createRainbowStore';
+import { $, AttachValue, attachValueSubscriptionMap, SignalFunction, Unsubscribe } from './signal';
+
+const ENABLE_LOGS = false;
+
+export const QueryStatuses = {
+  Idle: 'idle',
+  Loading: 'loading',
+  Success: 'success',
+  Error: 'error',
+} as const;
 
 /**
  * Represents the status of the remote data fetching process.
  */
-type RemoteStatus = 'idle' | 'loading' | 'success' | 'error';
+export type QueryStatus = (typeof QueryStatuses)[keyof typeof QueryStatuses];
 
 /**
  * Configuration options for remote data fetching.
@@ -26,6 +36,19 @@ interface CacheEntry<TData> {
   lastFetchedAt: number;
 }
 
+/**
+ * A specialized store interface combining Zustand's store API with remote fetching.
+ */
+export interface QueryStore<TData, TParams extends Record<string, unknown>, S extends StoreState<TData, TParams>>
+  extends UseBoundStore<StoreApi<S>> {
+  enabled: boolean;
+  destroy: () => void;
+  fetch: (params?: TParams, options?: FetchOptions) => Promise<void>;
+  isDataExpired: (override?: number) => boolean;
+  isStale: (override?: number) => boolean;
+  reset: () => void;
+}
+
 /**
  * The base store state including query-related fields and actions.
  */
@@ -35,7 +58,7 @@ type StoreState<TData, TParams extends Record<string, unknown>> = {
   error: Error | null;
   lastFetchedAt: number | null;
   queryCache: Record<string, CacheEntry<TData>>;
-  status: RemoteStatus;
+  status: QueryStatus;
   subscriptionCount: number;
   fetch: (params?: TParams, options?: FetchOptions) => Promise<void>;
   isDataExpired: (cacheTimeOverride?: number) => boolean;
@@ -43,23 +66,6 @@ type StoreState<TData, TParams extends Record<string, unknown>> = {
   reset: () => void;
 };
 
-/**
- * The keys that make up the internal state of the store.
- */
-type InternalStateKey = keyof StoreState<unknown, Record<string, unknown>>;
-
-/**
- * A specialized store interface combining Zustand's store API with remote fetching.
- */
-export interface QueryStore<TData, TParams extends Record<string, unknown>, S extends StoreState<TData, TParams>>
-  extends UseBoundStore<StoreApi<S>> {
-  enabled: boolean;
-  fetch: (params?: TParams, options?: FetchOptions) => Promise<void>;
-  isDataExpired: (override?: number) => boolean;
-  isStale: (override?: number) => boolean;
-  reset: () => void;
-}
-
 /**
  * Configuration options for creating a remote-enabled Rainbow store.
  */
@@ -68,17 +74,32 @@ type RainbowQueryStoreConfig<TQueryFnData, TParams extends Record<string, unknow
   onFetched?: (data: TData, store: QueryStore<TData, TParams, S>) => void;
   transform?: (data: TQueryFnData) => TData;
   cacheTime?: number;
-  defaultParams?: TParams;
+  params?: {
+    [K in keyof TParams]: ParamResolvable<TParams[K]>;
+  };
   disableDataCache?: boolean;
   enabled?: boolean;
-  queryKey: readonly unknown[] | ((params: TParams) => readonly unknown[]);
   staleTime?: number;
 };
 
-const SEVEN_DAYS = 1000 * 60 * 60 * 24 * 7;
-const TWO_MINUTES = 1000 * 60 * 2;
-const FIVE_SECONDS = 1000 * 5;
-const MIN_STALE_TIME = FIVE_SECONDS;
+/**
+ * A function that resolves to a value or an AttachValue wrapper.
+ */
+type ParamResolvable<T> = T | ((resolve: SignalFunction) => AttachValue<T>);
+
+/**
+ * The result of resolving parameters into their direct values and AttachValue wrappers.
+ */
+interface ResolvedParamsResult<TParams> {
+  directValues: Partial<TParams>;
+  paramAttachVals: Partial<Record<keyof TParams, AttachValue<unknown>>>;
+  resolvedParams: TParams;
+}
+
+/**
+ * The keys that make up the internal state of the store.
+ */
+type InternalStateKeys = keyof StoreState<unknown, Record<string, unknown>>;
 
 const [persist, discard] = [true, false];
 
@@ -97,22 +118,24 @@ const SHOULD_PERSIST_INTERNAL_STATE_MAP: Record<string, boolean> = {
   isStale: discard,
   reset: discard,
   subscriptionCount: discard,
-} satisfies Record<InternalStateKey, boolean>;
+} satisfies Record<InternalStateKeys, boolean>;
+
+const SEVEN_DAYS = 1000 * 60 * 60 * 24 * 7;
+const TWO_MINUTES = 1000 * 60 * 2;
+const FIVE_SECONDS = 1000 * 5;
+const MIN_STALE_TIME = FIVE_SECONDS;
 
-/**
- * Creates a remote-enabled Rainbow store with data fetching capabilities.
- * @template TQueryFnData - The raw data type returned by the fetcher
- * @template TParams - Parameters passed to the fetcher function
- * @template TData - The transformed data type (defaults to TQueryFnData)
- * @template U - Additional user-defined state
- */
 export function createRainbowQueryStore<
   TQueryFnData,
   TParams extends Record<string, unknown> = Record<string, unknown>,
   U = unknown,
   TData = TQueryFnData,
 >(
-  config: RainbowQueryStoreConfig<TQueryFnData, TParams, TData, StoreState<TData, TParams> & U>,
+  config: RainbowQueryStoreConfig<TQueryFnData, TParams, TData, StoreState<TData, TParams> & U> & {
+    params?: {
+      [K in keyof TParams]: ParamResolvable<TParams[K]>;
+    };
+  },
   customStateCreator?: StateCreator<U, [], [['zustand/subscribeWithSelector', never]]>,
   persistConfig?: RainbowPersistConfig<StoreState<TData, TParams> & U>
 ): QueryStore<TData, TParams, StoreState<TData, TParams> & U> {
@@ -123,16 +146,26 @@ export function createRainbowQueryStore<
     onFetched,
     transform,
     cacheTime = SEVEN_DAYS,
-    defaultParams,
+    params,
     disableDataCache = true,
     enabled = true,
-    queryKey,
     staleTime = TWO_MINUTES,
   } = config;
 
+  let paramAttachVals: Partial<Record<keyof TParams, AttachValue<unknown>>> = {};
+  let directValues: Partial<TParams> = {};
+
+  if (params) {
+    const result = resolveParams<TParams>(params);
+    paramAttachVals = result.paramAttachVals;
+    directValues = result.directValues;
+  }
+
   if (IS_DEV && staleTime < MIN_STALE_TIME) {
     console.warn(
-      `[RainbowQueryStore${persistConfig?.storageKey ? `: ${persistConfig.storageKey}` : ''}] ❌ Stale times under ${MIN_STALE_TIME / 1000} seconds are not recommended.`
+      `[RainbowQueryStore${persistConfig?.storageKey ? `: ${persistConfig.storageKey}` : ''}] ❌ Stale times under ${
+        MIN_STALE_TIME / 1000
+      } seconds are not recommended.`
     );
   }
 
@@ -140,11 +173,6 @@ export function createRainbowQueryStore<
   let activeRefetchTimeout: NodeJS.Timeout | null = null;
   let lastFetchKey: string | null = null;
 
-  const getQueryKey = (params: TParams): string => {
-    const key = typeof queryKey === 'function' ? queryKey(params) : queryKey;
-    return JSON.stringify(key);
-  };
-
   const initialData = {
     data: null,
     enabled,
@@ -155,9 +183,32 @@ export function createRainbowQueryStore<
     subscriptionCount: 0,
   };
 
-  const createState: StateCreator<S, [], [['zustand/subscribeWithSelector', never]]> = (set, get, api) => {
-    let isRefetchScheduled = false;
+  const getQueryKey = (params: TParams): string => JSON.stringify(Object.values(params));
 
+  const getCurrentResolvedParams = () => {
+    const currentParams = { ...directValues };
+    for (const k in paramAttachVals) {
+      const attachVal = paramAttachVals[k as keyof TParams];
+      if (!attachVal) continue;
+      currentParams[k as keyof TParams] = attachVal.value as TParams[keyof TParams];
+    }
+    return currentParams as TParams;
+  };
+
+  const scheduleNextFetch = (params: TParams) => {
+    if (staleTime <= 0) return;
+    if (activeRefetchTimeout) {
+      clearTimeout(activeRefetchTimeout);
+      activeRefetchTimeout = null;
+    }
+    activeRefetchTimeout = setTimeout(() => {
+      if (baseStore.getState().subscriptionCount > 0) {
+        baseStore.getState().fetch(params, { force: true });
+      }
+    }, staleTime);
+  };
+
+  const createState: StateCreator<S, [], [['zustand/subscribeWithSelector', never]]> = (set, get, api) => {
     const pruneCache = (state: S): S => {
       if (disableDataCache) return state;
       const now = Date.now();
@@ -170,27 +221,14 @@ export function createRainbowQueryStore<
       return { ...state, queryCache: newCache };
     };
 
-    const scheduleNextFetch = (params: TParams) => {
-      if (isRefetchScheduled || staleTime <= 0) return;
-      if (activeRefetchTimeout) clearTimeout(activeRefetchTimeout);
-
-      isRefetchScheduled = true;
-      activeRefetchTimeout = setTimeout(() => {
-        isRefetchScheduled = false;
-        if (get().subscriptionCount > 0) {
-          get().fetch(params, { force: true });
-        }
-      }, staleTime);
-    };
-
     const baseMethods = {
       async fetch(params: TParams | undefined, options: FetchOptions | undefined) {
         if (!get().enabled) return;
-
-        const effectiveParams = params ?? defaultParams ?? ({} as TParams);
+        const effectiveParams = params ?? getCurrentResolvedParams();
         const currentQueryKey = getQueryKey(effectiveParams);
+        const isLoading = get().status === 'loading';
 
-        if (activeFetchPromise && lastFetchKey === currentQueryKey && get().status === 'loading' && !options?.force) {
+        if (activeFetchPromise && lastFetchKey === currentQueryKey && isLoading && !options?.force) {
           return activeFetchPromise;
         }
 
@@ -207,9 +245,15 @@ export function createRainbowQueryStore<
 
         const fetchOperation = async () => {
           try {
-            const result = await fetcher(effectiveParams);
-            const rawData = result instanceof Promise ? await result : result;
-            const transformedData = transform ? transform(rawData) : (rawData as TData);
+            const rawResult = await fetcher(effectiveParams);
+            let transformedData: TData;
+            try {
+              transformedData = transform ? transform(rawResult) : (rawResult as TData);
+            } catch (transformError) {
+              throw new RainbowError(`[createRainbowQueryStore: ${persistConfig?.storageKey || currentQueryKey}]: transform failed`, {
+                cause: transformError,
+              });
+            }
 
             set(state => {
               const newState = {
@@ -237,15 +281,23 @@ export function createRainbowQueryStore<
             scheduleNextFetch(effectiveParams);
 
             if (onFetched) {
-              onFetched(transformedData, queryCapableStore);
+              try {
+                onFetched(transformedData, queryCapableStore);
+              } catch (onFetchedError) {
+                logger.error(
+                  new RainbowError(
+                    `[createRainbowQueryStore: ${persistConfig?.storageKey || currentQueryKey}]: onFetched callback failed`,
+                    { cause: onFetchedError }
+                  )
+                );
+              }
             }
           } catch (error) {
             logger.error(
               new RainbowError(`[createRainbowQueryStore: ${persistConfig?.storageKey || currentQueryKey}]: Failed to fetch data`),
               { error }
             );
-            // TODO: Improve retry logic
-            set(state => ({ ...state, error, status: 'error' as const }));
+            set(state => ({ ...state, error: error as Error, status: 'error' as const }));
             scheduleNextFetch(effectiveParams);
           } finally {
             activeFetchPromise = null;
@@ -278,12 +330,10 @@ export function createRainbowQueryStore<
         }
         activeFetchPromise = null;
         lastFetchKey = null;
-        isRefetchScheduled = false;
         set(state => ({ ...state, ...initialData }));
       },
     };
 
-    /* If customStateCreator is provided, it will return user-defined fields (U) */
     const userState = customStateCreator?.(set, get, api) ?? ({} as U);
 
     const subscribeWithSelector = api.subscribe;
@@ -294,17 +344,19 @@ export function createRainbowQueryStore<
       const handleSetEnabled = subscribeWithSelector((state: S, prev: S) => {
         if (state.enabled !== prev.enabled) {
           if (state.enabled) {
-            if (!state.data || state.isStale()) {
-              state.fetch(defaultParams);
+            const currentKey = getQueryKey(getCurrentResolvedParams());
+            if (currentKey !== lastFetchKey) {
+              state.fetch(getCurrentResolvedParams(), { force: true });
+            } else if (!state.data || state.isStale()) {
+              state.fetch();
             } else {
-              scheduleNextFetch(defaultParams ?? ({} as TParams));
+              scheduleNextFetch(getCurrentResolvedParams());
             }
           } else {
             if (activeRefetchTimeout) {
               clearTimeout(activeRefetchTimeout);
               activeRefetchTimeout = null;
             }
-            isRefetchScheduled = false;
           }
         }
       });
@@ -312,9 +364,9 @@ export function createRainbowQueryStore<
       const { data, fetch, isStale } = get();
 
       if (!data || isStale()) {
-        fetch(defaultParams, { force: true });
+        fetch(getCurrentResolvedParams(), { force: true });
       } else {
-        scheduleNextFetch(defaultParams ?? ({} as TParams));
+        scheduleNextFetch(getCurrentResolvedParams());
       }
 
       return () => {
@@ -327,7 +379,6 @@ export function createRainbowQueryStore<
               clearTimeout(activeRefetchTimeout);
               activeRefetchTimeout = null;
             }
-            isRefetchScheduled = false;
           }
           return { ...prev, subscriptionCount: newCount };
         });
@@ -359,30 +410,88 @@ export function createRainbowQueryStore<
     isStale: (override?: number) => baseStore.getState().isStale(override),
     reset: () => baseStore.getState().reset(),
     enabled,
+    destroy: () => {
+      for (const unsub of paramUnsubscribes) {
+        unsub();
+      }
+      paramUnsubscribes.length = 0;
+      queryCapableStore.getState().reset();
+    },
   });
 
+  const onParamChange = () => {
+    const newParams = getCurrentResolvedParams();
+    queryCapableStore.fetch(newParams, { force: true });
+  };
+
+  const paramUnsubscribes: Unsubscribe[] = [];
+
+  for (const k in paramAttachVals) {
+    const attachVal = paramAttachVals[k];
+    if (!attachVal) continue;
+
+    const subscribeFn = attachValueSubscriptionMap.get(attachVal);
+    if (ENABLE_LOGS) console.log('[🌀 ParamSubscription 🌀] Subscribed to param:', k);
+
+    if (subscribeFn) {
+      let oldVal = attachVal.value;
+      const unsub = subscribeFn(() => {
+        const newVal = attachVal.value;
+        if (!Object.is(oldVal, newVal)) {
+          oldVal = newVal;
+          if (ENABLE_LOGS) console.log('[🌀 ParamChange 🌀] Param changed:', k);
+          onParamChange();
+        }
+      });
+      paramUnsubscribes.push(unsub);
+    }
+  }
+
   return queryCapableStore;
 }
 
-/**
- * Creates a combined partialize function that ensures internal query state is always
- * persisted while respecting user-defined persistence preferences.
- */
+function isParamFn<T>(param: T | ((resolve: SignalFunction) => AttachValue<T>)): param is (resolve: SignalFunction) => AttachValue<T> {
+  return typeof param === 'function';
+}
+
+function resolveParams<TParams extends Record<string, unknown>>(params: {
+  [K in keyof TParams]: ParamResolvable<TParams[K]>;
+}): ResolvedParamsResult<TParams> {
+  const resolvedParams = {} as TParams;
+  const paramAttachVals: Partial<Record<keyof TParams, AttachValue<unknown>>> = {};
+  const directValues: Partial<TParams> = {};
+
+  for (const key in params) {
+    const param = params[key];
+    if (isParamFn(param)) {
+      const attachVal = param($);
+      resolvedParams[key] = attachVal.value as TParams[typeof key];
+      paramAttachVals[key] = attachVal;
+    } else {
+      resolvedParams[key] = param as TParams[typeof key];
+      directValues[key] = param as TParams[typeof key];
+    }
+  }
+
+  return { resolvedParams, paramAttachVals, directValues };
+}
+
 function createBlendedPartialize<TData, TParams extends Record<string, unknown>, S extends StoreState<TData, TParams> & U, U = unknown>(
   userPartialize: ((state: StoreState<TData, TParams> & U) => Partial<StoreState<TData, TParams> & U>) | undefined
 ) {
   return (state: S) => {
+    const clonedState = { ...state };
     const internalStateToPersist: Partial<S> = {};
 
-    for (const key in state) {
+    for (const key in clonedState) {
       if (key in SHOULD_PERSIST_INTERNAL_STATE_MAP) {
-        if (SHOULD_PERSIST_INTERNAL_STATE_MAP[key]) internalStateToPersist[key] = state[key];
-        delete state[key];
+        if (SHOULD_PERSIST_INTERNAL_STATE_MAP[key]) internalStateToPersist[key] = clonedState[key];
+        delete clonedState[key];
       }
     }
 
     return {
-      ...(userPartialize ? userPartialize(state) : state),
+      ...(userPartialize ? userPartialize(clonedState) : clonedState),
       ...internalStateToPersist,
     } satisfies Partial<S>;
   };
diff --git a/src/state/internal/createStore.ts b/src/state/internal/createStore.ts
index 3c49c5eb18a..05006491549 100644
--- a/src/state/internal/createStore.ts
+++ b/src/state/internal/createStore.ts
@@ -8,6 +8,9 @@ export type StoreWithPersist<TState> = Mutate<StoreApi<TState>, [['zustand/persi
   initializer: Initializer<TState>;
 };
 
+/**
+ * @deprecated This is a legacy store creator. Use `createRainbowStore` instead.
+ */
 export function createStore<TState>(
   initializer: Initializer<TState>,
   { persist: persistOptions }: { persist?: PersistOptions<TState> } = {}
diff --git a/src/state/internal/signal.ts b/src/state/internal/signal.ts
new file mode 100644
index 00000000000..d2d2b648b50
--- /dev/null
+++ b/src/state/internal/signal.ts
@@ -0,0 +1,165 @@
+import { StoreApi } from 'zustand';
+
+const ENABLE_LOGS = false;
+
+/* Store subscribe function so we can handle param changes on any attachVal (root or nested) */
+export const attachValueSubscriptionMap = new WeakMap<AttachValue<unknown>, Subscribe>();
+
+/* Global caching for top-level attachValues */
+const storeSignalCache = new WeakMap<
+  StoreApi<unknown>,
+  Map<(state: unknown) => unknown, Map<(a: unknown, b: unknown) => boolean, AttachValue<unknown>>>
+>();
+
+export type AttachValue<T> = T & { value: T } & {
+  readonly [K in keyof T]: AttachValue<T[K]>;
+};
+
+export type SignalFunction = {
+  <T>(store: StoreApi<T>): AttachValue<T>;
+  <T, S>(store: StoreApi<T>, selector: (state: T) => S, equalityFn?: (a: S, b: S) => boolean): AttachValue<S>;
+};
+
+export type Unsubscribe = () => void;
+export type Subscribe = (callback: () => void) => Unsubscribe;
+export type GetValue = () => unknown;
+export type SetValue = (path: unknown[], value: unknown) => void;
+
+const identity = <T>(x: T): T => x;
+
+const updateValue = <T>(obj: T, path: unknown[], value: unknown): T => {
+  if (!path.length) {
+    return value as T;
+  }
+  const [first, ...rest] = path;
+  const prevValue = (obj as Record<string, unknown>)[first as string];
+  const nextValue = updateValue(prevValue, rest, value);
+  if (Object.is(prevValue, nextValue)) {
+    return obj;
+  }
+  const copied = Array.isArray(obj) ? obj.slice() : { ...obj };
+  (copied as Record<string, unknown>)[first as string] = nextValue;
+  return copied as T;
+};
+
+export const createSignal = <T, S>(
+  store: StoreApi<T>,
+  selector: (state: T) => S,
+  equalityFn: (a: S, b: S) => boolean
+): [Subscribe, GetValue, SetValue] => {
+  let selected = selector(store.getState());
+  const listeners = new Set<() => void>();
+  let unsubscribe: Unsubscribe | undefined;
+
+  const sub: Subscribe = callback => {
+    if (!listeners.size) {
+      unsubscribe = store.subscribe(() => {
+        const nextSelected = selector(store.getState());
+        if (!equalityFn(selected, nextSelected)) {
+          selected = nextSelected;
+          listeners.forEach(listener => listener());
+        }
+      });
+    }
+    listeners.add(callback);
+    return () => {
+      listeners.delete(callback);
+      if (!listeners.size && unsubscribe) {
+        unsubscribe();
+        unsubscribe = undefined;
+      }
+    };
+  };
+
+  const get: GetValue = () => {
+    if (!listeners.size) {
+      selected = selector(store.getState());
+    }
+    return selected;
+  };
+
+  const set: SetValue = (path, value) => {
+    if (selector !== identity) {
+      throw new Error('Cannot set a value with a selector');
+    }
+    store.setState(prev => updateValue(prev, path, value), true);
+  };
+
+  return [sub, get, set];
+};
+
+function getOrCreateAttachValue<T, S>(store: StoreApi<T>, selector: (state: T) => S, equalityFn: (a: S, b: S) => boolean): AttachValue<S> {
+  let bySelector = storeSignalCache.get(store);
+  if (!bySelector) {
+    bySelector = new Map();
+    storeSignalCache.set(store, bySelector);
+  }
+
+  let byEqFn = bySelector.get(selector as (state: unknown) => unknown);
+  if (!byEqFn) {
+    byEqFn = new Map();
+    bySelector.set(selector as (state: unknown) => unknown, byEqFn);
+  }
+
+  const existing = byEqFn.get(equalityFn as (a: unknown, b: unknown) => boolean);
+  if (existing) {
+    return existing as AttachValue<S>;
+  }
+
+  const [subscribe, getVal, setVal] = createSignal(store, selector, equalityFn);
+
+  if (ENABLE_LOGS) console.log('[🌀 AttachValue 🌀] Created root attachValue:', { selector: selector.toString() });
+
+  const localCache = new Map<string, AttachValue<unknown>>();
+
+  const createAttachValue = (fullPath: string): AttachValue<unknown> => {
+    const handler: ProxyHandler<object> = {
+      get(_, key) {
+        if (key === 'value') {
+          let v = getVal();
+          const parts = fullPath.split('.');
+          for (const p of parts) {
+            if (p) v = (v as Record<string, unknown>)[p];
+          }
+          return v;
+        }
+        const pathKey = fullPath ? `${fullPath}.${key.toString()}` : key.toString();
+        const cached = localCache.get(pathKey);
+        if (cached) {
+          console.log('[🌀 AttachValue 🌀] Cache hit for:', pathKey);
+          return cached;
+        }
+        const val = createAttachValue(pathKey);
+        attachValueSubscriptionMap.set(val, subscribe);
+        localCache.set(pathKey, val);
+        return val;
+      },
+      set(_, __, value) {
+        const path = fullPath.split('.');
+        if (path[0] === '') path.shift();
+        setVal(path, value);
+        return true;
+      },
+    };
+
+    return new Proxy(Object.create(null), handler) as AttachValue<unknown>;
+  };
+
+  const rootVal = createAttachValue('');
+  subscribe(() => {
+    return;
+  });
+  attachValueSubscriptionMap.set(rootVal, subscribe);
+  byEqFn.set(equalityFn as (a: unknown, b: unknown) => boolean, rootVal);
+  return rootVal as AttachValue<S>;
+}
+
+export function $<T>(store: StoreApi<T>): AttachValue<T>;
+export function $<T, S>(store: StoreApi<T>, selector: (state: T) => S, equalityFn?: (a: S, b: S) => boolean): AttachValue<S>;
+export function $(
+  store: StoreApi<unknown>,
+  selector: (state: unknown) => unknown = identity,
+  equalityFn: (a: unknown, b: unknown) => boolean = Object.is
+) {
+  return getOrCreateAttachValue(store, selector, equalityFn);
+}
diff --git a/src/state/internal/tests/RainbowQueryStoreTest.tsx b/src/state/internal/tests/RainbowQueryStoreTest.tsx
index a3e22886ab2..b4444968cb9 100644
--- a/src/state/internal/tests/RainbowQueryStoreTest.tsx
+++ b/src/state/internal/tests/RainbowQueryStoreTest.tsx
@@ -7,97 +7,147 @@ import { SupportedCurrencyKey } from '@/references';
 import { queryUserAssets } from '@/__swaps__/screens/Swap/resources/assets/userAssets';
 import { ParsedAssetsDictByChain } from '@/__swaps__/types/assets';
 import { createRainbowQueryStore } from '../createRainbowQueryStore';
+import { createRainbowStore } from '../createRainbowStore';
 
-function getRandomAddress() {
-  return Math.random() < 0.5 ? '0x2e67869829c734ac13723A138a952F7A8B56e774' : '0xCFB83E14AEd465c79F3F82f4cfF4ff7965897644';
-}
+const ENABLE_LOGS = false;
 
-type QueryParams = { address: Address; currency: SupportedCurrencyKey };
+type AddressStore = {
+  address: Address;
+  currency: SupportedCurrencyKey;
+  nestedAddressTest: {
+    address: Address;
+  };
+  setAddress: (address: Address) => void;
+};
+
+const testAddresses: Address[] = [
+  '0x2e67869829c734ac13723A138a952F7A8B56e774',
+  '0xCFB83E14AEd465c79F3F82f4cfF4ff7965897644',
+  '0x17cd072cBd45031EFc21Da538c783E0ed3b25DCc',
+];
+
+const useAddressStore = createRainbowStore<AddressStore>((set, get) => ({
+  address: testAddresses[0],
+  currency: 'USD',
+  nestedAddressTest: { address: testAddresses[0] },
+
+  setAddress: (address: Address) => {
+    set({ address });
+    console.log('DID ADDRESS SET?', 'new address:', get().address);
+  },
+}));
 
 type TestStore = {
   userAssets: ParsedAssetsDictByChain;
   getHighestValueAsset: () => number;
   setUserAssets: (data: ParsedAssetsDictByChain) => void;
 };
+type QueryParams = { address: Address; currency: SupportedCurrencyKey };
+
+function logFetchInfo(params: QueryParams) {
+  console.log('[PARAMS]:', JSON.stringify(params, null, 2));
+  const formattedTimeWithSeconds = new Date(Date.now()).toLocaleTimeString('en-US', {
+    hour: '2-digit',
+    minute: '2-digit',
+    second: '2-digit',
+  });
+  if (ENABLE_LOGS) {
+    console.log('[🔄 Requesting Fetch] - Last Fetch Attempt:', formattedTimeWithSeconds, '\nParams:', {
+      address: params.address,
+      currency: params.currency,
+    });
+  }
+}
 
-export const userAssetsStore = createRainbowQueryStore<ParsedAssetsDictByChain, QueryParams, TestStore, ParsedAssetsDictByChain>(
+export const userAssetsTestStore = createRainbowQueryStore<ParsedAssetsDictByChain, QueryParams, TestStore>(
   {
-    fetcher: () => queryUserAssets({ address: getRandomAddress(), currency: 'USD' }),
-    // onFetched: (data, store) => store.setState({ userAssets: data }),
-    transform: data => {
-      const formattedTimeWithSeconds = new Date(Date.now()).toLocaleTimeString('en-US', {
-        hour: '2-digit',
-        minute: '2-digit',
-        second: '2-digit',
-      });
-      console.log('[Transform - Last Fetch Attempt]: ', formattedTimeWithSeconds);
-      return data;
+    fetcher: ({ address, currency }) => {
+      if (ENABLE_LOGS) logFetchInfo({ address, currency });
+      return queryUserAssets({ address, currency });
     },
+    onFetched: (data, store) => store.setState({ userAssets: data }),
 
-    disableDataCache: false,
-    queryKey: ['userAssets'],
-    staleTime: 30 * 60 * 1000, // 30m
+    params: {
+      address: $ => $(useAddressStore).address,
+      currency: $ => $(useAddressStore).currency,
+    },
+    staleTime: 20 * 1000, // 20s
   },
+
   (set, get) => ({
     userAssets: [],
 
-    getHighestValueAsset: () => {
-      const data = get().userAssets;
-      const highestValueAsset = Object.values(data)
+    getHighestValueAsset: () =>
+      Object.values(get().userAssets)
         .flatMap(chainAssets => Object.values(chainAssets))
-        .reduce((max, asset) => {
-          return Math.max(max, Number(asset.balance.display));
-        }, 0);
-      return highestValueAsset;
-    },
+        .reduce((max, asset) => Math.max(max, Number(asset.balance.display)), 0),
 
     setUserAssets: (data: ParsedAssetsDictByChain) => set({ userAssets: data }),
-  }),
-  {
-    // partialize: state => ({ userAssets: state.userAssets }),
-    storageKey: 'userAssetsTesting79876',
-  }
+  })
 );
 
 export const UserAssetsTest = memo(function UserAssetsTest() {
-  const data = userAssetsStore(state => state.data);
-  const enabled = userAssetsStore(state => state.enabled);
-
-  console.log('RERENDER - enabled:', enabled);
+  const data = userAssetsTestStore(state => state.userAssets);
+  const enabled = userAssetsTestStore(state => state.enabled);
 
   useEffect(() => {
-    const status = userAssetsStore.getState().status;
-    const isFetching = status === 'loading';
-    // eslint-disable-next-line no-nested-ternary
-    const emojiForStatus = isFetching ? '🔄' : status === 'success' ? '✅' : '❌';
-    console.log('[NEW STATUS]:', emojiForStatus, status);
-
-    if (data) {
+    if (ENABLE_LOGS) {
       const first5Tokens = Object.values(data)
         .flatMap(chainAssets => Object.values(chainAssets))
         .slice(0, 5);
-      console.log('[First 5 Token Symbols]:', first5Tokens.map(token => token.symbol).join(', '));
+      console.log('[🔔 UserAssetsTest 🔔] userAssets data updated - first 5 tokens:', first5Tokens.map(token => token.symbol).join(', '));
     }
   }, [data]);
 
+  useEffect(() => {
+    if (ENABLE_LOGS) console.log(`[🔔 UserAssetsTest 🔔] enabled updated to: ${enabled ? '✅ ENABLED' : '🛑 DISABLED'}`);
+  }, [enabled]);
+
   return (
     data && (
       <View style={styles.container}>
         <Text color="label" size="17pt" weight="heavy">
           Number of assets: {Object.values(data).reduce((acc, chainAssets) => acc + Object.keys(chainAssets).length, 0)}
         </Text>
-        <ButtonPressAnimation onPress={() => userAssetsStore.setState({ enabled: !enabled })} style={styles.button}>
-          <Text color="label" size="17pt" weight="heavy">
-            {enabled ? 'Disable fetching' : 'Enable fetching'}
-          </Text>
-        </ButtonPressAnimation>
+        <View style={styles.buttonGroup}>
+          <ButtonPressAnimation
+            onPress={() => {
+              const currentAddress = useAddressStore.getState().nestedAddressTest.address;
+              switch (currentAddress) {
+                case testAddresses[0]:
+                  useAddressStore.getState().setAddress(testAddresses[1]);
+                  break;
+                case testAddresses[1]:
+                  useAddressStore.getState().setAddress(testAddresses[2]);
+                  break;
+                case testAddresses[2]:
+                  useAddressStore.getState().setAddress(testAddresses[0]);
+                  break;
+              }
+            }}
+            style={styles.button}
+          >
+            <Text color="label" size="17pt" weight="heavy">
+              Shuffle Address
+            </Text>
+          </ButtonPressAnimation>
+          <ButtonPressAnimation
+            onPress={() => {
+              userAssetsTestStore.setState({ enabled: !enabled });
+            }}
+            style={styles.button}
+          >
+            <Text color="label" size="17pt" weight="heavy">
+              {userAssetsTestStore.getState().enabled ? 'Disable fetching' : 'Enable fetching'}
+            </Text>
+          </ButtonPressAnimation>
+        </View>
       </View>
     )
   );
 });
 
-const initialData = userAssetsStore.getState().data;
-console.log('[Initial Data Exists]:', !!initialData);
+if (ENABLE_LOGS) console.log('[💾 UserAssetsTest 💾] initial data exists:', !!userAssetsTestStore.getState().userAssets);
 
 const styles = StyleSheet.create({
   button: {
@@ -108,6 +158,12 @@ const styles = StyleSheet.create({
     justifyContent: 'center',
     paddingHorizontal: 20,
   },
+  buttonGroup: {
+    alignItems: 'center',
+    flexDirection: 'column',
+    justifyContent: 'center',
+    gap: 24,
+  },
   container: {
     alignItems: 'center',
     flex: 1,

From eef99038ce144c4835f4524f2ae6d6ce383d20cc Mon Sep 17 00:00:00 2001
From: Christian Baroni <7061887+christianbaroni@users.noreply.github.com>
Date: Mon, 16 Dec 2024 20:36:38 +0000
Subject: [PATCH 09/29] Misc. cleanup

---
 src/state/internal/createRainbowQueryStore.ts |  7 +++
 src/state/internal/createRainbowStore.ts      | 58 +++++++++----------
 .../internal/tests/RainbowQueryStoreTest.tsx  | 30 ++++++----
 3 files changed, 54 insertions(+), 41 deletions(-)

diff --git a/src/state/internal/createRainbowQueryStore.ts b/src/state/internal/createRainbowQueryStore.ts
index bb39bd07df7..a629cf54840 100644
--- a/src/state/internal/createRainbowQueryStore.ts
+++ b/src/state/internal/createRainbowQueryStore.ts
@@ -125,6 +125,13 @@ const TWO_MINUTES = 1000 * 60 * 2;
 const FIVE_SECONDS = 1000 * 5;
 const MIN_STALE_TIME = FIVE_SECONDS;
 
+/**
+ * Creates a query-enabled Rainbow store with data fetching capabilities.
+ * @template TQueryFnData - The raw data type returned by the fetcher
+ * @template TParams - Parameters passed to the fetcher function
+ * @template U - User-defined custom store state
+ * @template TData - The transformed data type, if applicable (defaults to `TQueryFnData`)
+ */
 export function createRainbowQueryStore<
   TQueryFnData,
   TParams extends Record<string, unknown> = Record<string, unknown>,
diff --git a/src/state/internal/createRainbowStore.ts b/src/state/internal/createRainbowStore.ts
index 371095b6c00..5de0748c95d 100644
--- a/src/state/internal/createRainbowStore.ts
+++ b/src/state/internal/createRainbowStore.ts
@@ -44,6 +44,35 @@ export interface RainbowPersistConfig<S> {
   version?: number;
 }
 
+/**
+ * Creates a Rainbow store with optional persistence functionality.
+ * @param createState - The state creator function for the Rainbow store.
+ * @param persistConfig - The configuration options for the persistable Rainbow store.
+ * @returns A Zustand store with the specified state and optional persistence.
+ */
+export function createRainbowStore<S = unknown>(
+  createState: StateCreator<S, [], [['zustand/subscribeWithSelector', never]]>,
+  persistConfig?: RainbowPersistConfig<S>
+) {
+  if (!persistConfig) {
+    return create<S>()(subscribeWithSelector(createState));
+  }
+
+  const { persistStorage, version } = createPersistStorage(persistConfig);
+
+  return create<S>()(
+    subscribeWithSelector(
+      persist(createState, {
+        migrate: persistConfig.migrate,
+        name: persistConfig.storageKey,
+        partialize: persistConfig.partialize || (state => state),
+        storage: persistStorage,
+        version,
+      })
+    )
+  );
+}
+
 /**
  * Creates a persist storage object for the Rainbow store.
  * @param config - The configuration options for the persistable Rainbow store.
@@ -131,32 +160,3 @@ function defaultDeserializeState<S>(serializedState: string): StorageValue<Parti
     throw error;
   }
 }
-
-/**
- * Creates a Rainbow store with optional persistence functionality.
- * @param createState - The state creator function for the Rainbow store.
- * @param persistConfig - The configuration options for the persistable Rainbow store.
- * @returns A Zustand store with the specified state and optional persistence.
- */
-export function createRainbowStore<S = unknown>(
-  createState: StateCreator<S, [], [['zustand/subscribeWithSelector', never]]>,
-  persistConfig?: RainbowPersistConfig<S>
-) {
-  if (!persistConfig) {
-    return create<S>()(subscribeWithSelector(createState));
-  }
-
-  const { persistStorage, version } = createPersistStorage(persistConfig);
-
-  return create<S>()(
-    subscribeWithSelector(
-      persist(createState, {
-        migrate: persistConfig.migrate,
-        name: persistConfig.storageKey,
-        partialize: persistConfig.partialize || (state => state),
-        storage: persistStorage,
-        version,
-      })
-    )
-  );
-}
diff --git a/src/state/internal/tests/RainbowQueryStoreTest.tsx b/src/state/internal/tests/RainbowQueryStoreTest.tsx
index b4444968cb9..4c18ff44138 100644
--- a/src/state/internal/tests/RainbowQueryStoreTest.tsx
+++ b/src/state/internal/tests/RainbowQueryStoreTest.tsx
@@ -33,7 +33,7 @@ const useAddressStore = createRainbowStore<AddressStore>((set, get) => ({
 
   setAddress: (address: Address) => {
     set({ address });
-    console.log('DID ADDRESS SET?', 'new address:', get().address);
+    if (ENABLE_LOGS) console.log('[👤 useAddressStore 👤] New address set:', get().address);
   },
 }));
 
@@ -45,20 +45,26 @@ type TestStore = {
 type QueryParams = { address: Address; currency: SupportedCurrencyKey };
 
 function logFetchInfo(params: QueryParams) {
-  console.log('[PARAMS]:', JSON.stringify(params, null, 2));
+  console.log('[🔄 logFetchInfo 🔄] Current params:', JSON.stringify(Object.values(params), null, 2));
   const formattedTimeWithSeconds = new Date(Date.now()).toLocaleTimeString('en-US', {
     hour: '2-digit',
     minute: '2-digit',
     second: '2-digit',
   });
-  if (ENABLE_LOGS) {
-    console.log('[🔄 Requesting Fetch] - Last Fetch Attempt:', formattedTimeWithSeconds, '\nParams:', {
-      address: params.address,
-      currency: params.currency,
-    });
-  }
+  console.log('[🔄 Requesting Fetch 🔄] Last fetch attempt:', formattedTimeWithSeconds, '\nParams:', {
+    address: params.address,
+    currency: params.currency,
+  });
 }
 
+const time = {
+  seconds: (n: number) => n * 1000,
+  minutes: (n: number) => time.seconds(n * 60),
+  hours: (n: number) => time.minutes(n * 60),
+  days: (n: number) => time.hours(n * 24),
+  weeks: (n: number) => time.days(n * 7),
+};
+
 export const userAssetsTestStore = createRainbowQueryStore<ParsedAssetsDictByChain, QueryParams, TestStore>(
   {
     fetcher: ({ address, currency }) => {
@@ -71,7 +77,7 @@ export const userAssetsTestStore = createRainbowQueryStore<ParsedAssetsDictByCha
       address: $ => $(useAddressStore).address,
       currency: $ => $(useAddressStore).currency,
     },
-    staleTime: 20 * 1000, // 20s
+    staleTime: time.minutes(1),
   },
 
   (set, get) => ({
@@ -138,7 +144,7 @@ export const UserAssetsTest = memo(function UserAssetsTest() {
             style={styles.button}
           >
             <Text color="label" size="17pt" weight="heavy">
-              {userAssetsTestStore.getState().enabled ? 'Disable fetching' : 'Enable fetching'}
+              {userAssetsTestStore.getState().enabled ? 'Disable Fetching' : 'Enable Fetching'}
             </Text>
           </ButtonPressAnimation>
         </View>
@@ -147,7 +153,7 @@ export const UserAssetsTest = memo(function UserAssetsTest() {
   );
 });
 
-if (ENABLE_LOGS) console.log('[💾 UserAssetsTest 💾] initial data exists:', !!userAssetsTestStore.getState().userAssets);
+if (ENABLE_LOGS) console.log('[💾 UserAssetsTest 💾] Initial data exists:', !!userAssetsTestStore.getState().userAssets);
 
 const styles = StyleSheet.create({
   button: {
@@ -161,8 +167,8 @@ const styles = StyleSheet.create({
   buttonGroup: {
     alignItems: 'center',
     flexDirection: 'column',
-    justifyContent: 'center',
     gap: 24,
+    justifyContent: 'center',
   },
   container: {
     alignItems: 'center',

From 96d8fd40295d59077a5f678e5e9cb62b2b04f589 Mon Sep 17 00:00:00 2001
From: Christian Baroni <7061887+christianbaroni@users.noreply.github.com>
Date: Mon, 16 Dec 2024 20:42:21 +0000
Subject: [PATCH 10/29] Add time utility

---
 src/state/internal/createRainbowQueryStore.ts | 43 +++++++++++--------
 .../internal/tests/RainbowQueryStoreTest.tsx  | 10 +----
 2 files changed, 25 insertions(+), 28 deletions(-)

diff --git a/src/state/internal/createRainbowQueryStore.ts b/src/state/internal/createRainbowQueryStore.ts
index a629cf54840..eac8dd1ddd4 100644
--- a/src/state/internal/createRainbowQueryStore.ts
+++ b/src/state/internal/createRainbowQueryStore.ts
@@ -74,11 +74,11 @@ type RainbowQueryStoreConfig<TQueryFnData, TParams extends Record<string, unknow
   onFetched?: (data: TData, store: QueryStore<TData, TParams, S>) => void;
   transform?: (data: TQueryFnData) => TData;
   cacheTime?: number;
+  disableDataCache?: boolean;
+  enabled?: boolean;
   params?: {
     [K in keyof TParams]: ParamResolvable<TParams[K]>;
   };
-  disableDataCache?: boolean;
-  enabled?: boolean;
   staleTime?: number;
 };
 
@@ -120,10 +120,15 @@ const SHOULD_PERSIST_INTERNAL_STATE_MAP: Record<string, boolean> = {
   subscriptionCount: discard,
 } satisfies Record<InternalStateKeys, boolean>;
 
-const SEVEN_DAYS = 1000 * 60 * 60 * 24 * 7;
-const TWO_MINUTES = 1000 * 60 * 2;
-const FIVE_SECONDS = 1000 * 5;
-const MIN_STALE_TIME = FIVE_SECONDS;
+export const time = {
+  seconds: (n: number) => n * 1000,
+  minutes: (n: number) => time.seconds(n * 60),
+  hours: (n: number) => time.minutes(n * 60),
+  days: (n: number) => time.hours(n * 24),
+  weeks: (n: number) => time.days(n * 7),
+};
+
+const MIN_STALE_TIME = time.seconds(5);
 
 /**
  * Creates a query-enabled Rainbow store with data fetching capabilities.
@@ -152,22 +157,13 @@ export function createRainbowQueryStore<
     fetcher,
     onFetched,
     transform,
-    cacheTime = SEVEN_DAYS,
-    params,
+    cacheTime = time.days(7),
     disableDataCache = true,
     enabled = true,
-    staleTime = TWO_MINUTES,
+    params,
+    staleTime = time.minutes(2),
   } = config;
 
-  let paramAttachVals: Partial<Record<keyof TParams, AttachValue<unknown>>> = {};
-  let directValues: Partial<TParams> = {};
-
-  if (params) {
-    const result = resolveParams<TParams>(params);
-    paramAttachVals = result.paramAttachVals;
-    directValues = result.directValues;
-  }
-
   if (IS_DEV && staleTime < MIN_STALE_TIME) {
     console.warn(
       `[RainbowQueryStore${persistConfig?.storageKey ? `: ${persistConfig.storageKey}` : ''}] ❌ Stale times under ${
@@ -176,6 +172,15 @@ export function createRainbowQueryStore<
     );
   }
 
+  let directValues: Partial<TParams> = {};
+  let paramAttachVals: Partial<Record<keyof TParams, AttachValue<unknown>>> = {};
+
+  if (params) {
+    const result = resolveParams<TParams>(params);
+    paramAttachVals = result.paramAttachVals;
+    directValues = result.directValues;
+  }
+
   let activeFetchPromise: Promise<void> | null = null;
   let activeRefetchTimeout: NodeJS.Timeout | null = null;
   let lastFetchKey: string | null = null;
@@ -392,7 +397,7 @@ export function createRainbowQueryStore<
       };
     };
 
-    // Merge base data, user state, and methods into the final store state
+    /* Merge base data, user state, and methods into the final store state */
     return {
       ...initialData,
       ...userState,
diff --git a/src/state/internal/tests/RainbowQueryStoreTest.tsx b/src/state/internal/tests/RainbowQueryStoreTest.tsx
index 4c18ff44138..43fba79f800 100644
--- a/src/state/internal/tests/RainbowQueryStoreTest.tsx
+++ b/src/state/internal/tests/RainbowQueryStoreTest.tsx
@@ -6,7 +6,7 @@ import { Text } from '@/design-system';
 import { SupportedCurrencyKey } from '@/references';
 import { queryUserAssets } from '@/__swaps__/screens/Swap/resources/assets/userAssets';
 import { ParsedAssetsDictByChain } from '@/__swaps__/types/assets';
-import { createRainbowQueryStore } from '../createRainbowQueryStore';
+import { createRainbowQueryStore, time } from '../createRainbowQueryStore';
 import { createRainbowStore } from '../createRainbowStore';
 
 const ENABLE_LOGS = false;
@@ -57,14 +57,6 @@ function logFetchInfo(params: QueryParams) {
   });
 }
 
-const time = {
-  seconds: (n: number) => n * 1000,
-  minutes: (n: number) => time.seconds(n * 60),
-  hours: (n: number) => time.minutes(n * 60),
-  days: (n: number) => time.hours(n * 24),
-  weeks: (n: number) => time.days(n * 7),
-};
-
 export const userAssetsTestStore = createRainbowQueryStore<ParsedAssetsDictByChain, QueryParams, TestStore>(
   {
     fetcher: ({ address, currency }) => {

From 0b59554536b9e9e44dcf6c3bdd6b272fbf2a2835 Mon Sep 17 00:00:00 2001
From: Christian Baroni <7061887+christianbaroni@users.noreply.github.com>
Date: Mon, 16 Dec 2024 21:10:21 +0000
Subject: [PATCH 11/29] Guard console.log

---
 src/state/internal/signal.ts | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/state/internal/signal.ts b/src/state/internal/signal.ts
index d2d2b648b50..cea47c31da1 100644
--- a/src/state/internal/signal.ts
+++ b/src/state/internal/signal.ts
@@ -126,7 +126,7 @@ function getOrCreateAttachValue<T, S>(store: StoreApi<T>, selector: (state: T) =
         const pathKey = fullPath ? `${fullPath}.${key.toString()}` : key.toString();
         const cached = localCache.get(pathKey);
         if (cached) {
-          console.log('[🌀 AttachValue 🌀] Cache hit for:', pathKey);
+          if (ENABLE_LOGS) console.log('[🌀 AttachValue 🌀] Cache hit for:', pathKey);
           return cached;
         }
         const val = createAttachValue(pathKey);

From c02515d75b2c98e7b47e6ae35b682736890c4bf2 Mon Sep 17 00:00:00 2001
From: Christian Baroni <7061887+christianbaroni@users.noreply.github.com>
Date: Tue, 17 Dec 2024 21:49:21 +0000
Subject: [PATCH 12/29] Add setData, clean up caching and staleTime logic,
 rename to createQueryStore

---
 ...inbowQueryStore.ts => createQueryStore.ts} | 287 +++++++++++++-----
 ...wQueryStoreTest.tsx => QueryStoreTest.tsx} |  18 +-
 2 files changed, 224 insertions(+), 81 deletions(-)
 rename src/state/internal/{createRainbowQueryStore.ts => createQueryStore.ts} (58%)
 rename src/state/internal/tests/{RainbowQueryStoreTest.tsx => QueryStoreTest.tsx} (91%)

diff --git a/src/state/internal/createRainbowQueryStore.ts b/src/state/internal/createQueryStore.ts
similarity index 58%
rename from src/state/internal/createRainbowQueryStore.ts
rename to src/state/internal/createQueryStore.ts
index eac8dd1ddd4..4299e2a8fda 100644
--- a/src/state/internal/createRainbowQueryStore.ts
+++ b/src/state/internal/createQueryStore.ts
@@ -7,92 +7,212 @@ import { $, AttachValue, attachValueSubscriptionMap, SignalFunction, Unsubscribe
 
 const ENABLE_LOGS = false;
 
+/**
+ * A set of constants representing the various stages of a query's remote data fetching process.
+ */
 export const QueryStatuses = {
+  Error: 'error',
   Idle: 'idle',
   Loading: 'loading',
   Success: 'success',
-  Error: 'error',
 } as const;
 
 /**
- * Represents the status of the remote data fetching process.
+ * Represents the current status of the query's remote data fetching operation.
+ *
+ * Possible values:
+ * - **`'error'`**  : The most recent request encountered an error.
+ * - **`'idle'`**   : No request in progress, no error, no data yet.
+ * - **`'loading'`** : A request is currently in progress.
+ * - **`'success'`** : The most recent request has succeeded, and `data` is available.
  */
 export type QueryStatus = (typeof QueryStatuses)[keyof typeof QueryStatuses];
 
 /**
- * Configuration options for remote data fetching.
+ * Defines additional options for a data fetch operation.
  */
 interface FetchOptions {
+  /**
+   * Overrides the default cache duration for this fetch, in milliseconds.
+   * When data in the cache is older than this duration, it will be considered expired and
+   * become eligible for pruning.
+   */
   cacheTime?: number;
+  /**
+   * Forces a fetch request even if the current data is fresh and not stale.
+   * If `true`, the fetch operation bypasses existing cached data.
+   */
   force?: boolean;
+  /**
+   * Overrides the default stale duration for this fetch, in milliseconds.
+   * When data is older than this duration, it is considered stale and if the query is active,
+   * a background refetch will occur.
+   */
   staleTime?: number;
 }
 
 /**
- * Represents a cached query result.
+ * Represents an entry in the query cache, which stores fetched data along with metadata.
  */
 interface CacheEntry<TData> {
-  data: TData;
+  data: TData | null;
   lastFetchedAt: number;
 }
 
 /**
- * A specialized store interface combining Zustand's store API with remote fetching.
+ * A specialized store interface that combines Zustand's store capabilities with remote data fetching support.
+ *
+ * In addition to Zustand's store API (such as `getState()` and `subscribe()`), this interface provides:
+ * - **`enabled`**: A boolean indicating if the store is actively fetching data.
+ * - **`fetch(params, options)`**: Initiates a data fetch operation.
+ * - **`isDataExpired(override?)`**: Checks if the current data has expired based on `cacheTime`.
+ * - **`isStale(override?)`**: Checks if the current data is stale based on `staleTime`.
+ * - **`reset()`**: Resets the store to its initial state, clearing data and errors.
  */
 export interface QueryStore<TData, TParams extends Record<string, unknown>, S extends StoreState<TData, TParams>>
   extends UseBoundStore<StoreApi<S>> {
+  /**
+   * Indicates whether the store should actively fetch data.
+   * When `false`, the store won't automatically refetch data.
+   */
   enabled: boolean;
-  destroy: () => void;
+  /**
+   * Initiates a data fetch for the given parameters. If no parameters are provided, the store's
+   * current parameters are used.
+   * @param params - Optional parameters to pass to the fetcher function.
+   * @param options - Optional {@link FetchOptions} to customize the fetch behavior.
+   * @returns A promise that resolves when the fetch operation completes.
+   */
   fetch: (params?: TParams, options?: FetchOptions) => Promise<void>;
+  /**
+   * Returns the cached data, if available, for the current query params.
+   * @returns The cached data, or `null` if no data is available.
+   */
+  getData: (params?: TParams) => TData | null;
+  /**
+   * Determines if the current data is expired, meaning it has exceeded the `cacheTime` duration.
+   * @param override - An optional override for the default cache time, in milliseconds.
+   * @returns `true` if the data is expired, otherwise `false`.
+   */
   isDataExpired: (override?: number) => boolean;
+  /**
+   * Determines if the current data is stale, meaning it has exceeded the `staleTime` duration.
+   * Stale data may be refreshed automatically in the background.
+   * @param override - An optional override for the default stale time, in milliseconds.
+   * @returns `true` if the data is stale, otherwise `false`.
+   */
   isStale: (override?: number) => boolean;
+  /**
+   * Resets the store to its initial state, clearing data, error, and any cached values.
+   */
   reset: () => void;
 }
 
 /**
- * The base store state including query-related fields and actions.
+ * The state structure managed by the query store, including query-related fields and actions.
+ * This type is generally internal, and extended by user-defined states when creating a store.
  */
 type StoreState<TData, TParams extends Record<string, unknown>> = {
-  data: TData | null;
   enabled: boolean;
   error: Error | null;
   lastFetchedAt: number | null;
-  queryCache: Record<string, CacheEntry<TData>>;
+  queryCache: Record<string, CacheEntry<TData> | undefined>;
   status: QueryStatus;
   subscriptionCount: number;
   fetch: (params?: TParams, options?: FetchOptions) => Promise<void>;
+  getData: (params?: TParams) => TData | null;
   isDataExpired: (cacheTimeOverride?: number) => boolean;
   isStale: (staleTimeOverride?: number) => boolean;
   reset: () => void;
 };
 
 /**
- * Configuration options for creating a remote-enabled Rainbow store.
+ * Configuration options for creating a query-enabled Rainbow store.
  */
-type RainbowQueryStoreConfig<TQueryFnData, TParams extends Record<string, unknown>, TData, S extends StoreState<TData, TParams>> = {
+export type RainbowQueryStoreConfig<TQueryFnData, TParams extends Record<string, unknown>, TData, S extends StoreState<TData, TParams>> = {
+  /**
+   * A function responsible for fetching data from a remote source.
+   * Receives parameters of type `TParams` and returns either a promise or a raw data value of type `TQueryFnData`.
+   */
   fetcher: (params: TParams) => TQueryFnData | Promise<TQueryFnData>;
+  /**
+   * A callback invoked whenever fresh data is successfully fetched.
+   * Receives the transformed data and the store instance, allowing for side effects or additional updates.
+   */
   onFetched?: (data: TData, store: QueryStore<TData, TParams, S>) => void;
+  /**
+   * A function that overrides the default behavior of setting the fetched data in the store's query cache.
+   * Receives the transformed data and a set function that updates the store state.
+   */
+  setData?: (data: TData, set: (partial: S | Partial<S> | ((state: S) => S | Partial<S>)) => void) => void;
+  /**
+   * Suppresses warnings in the event a `staleTime` under the minimum is desired.
+   * @default false
+   */
+  suppressStaleTimeWarning?: boolean;
+  /**
+   * A function to transform the raw fetched data (`TQueryFnData`) into another form (`TData`).
+   * If not provided, the raw data returned by `fetcher` is used.
+   */
   transform?: (data: TQueryFnData) => TData;
+  /**
+   * The maximum duration, in milliseconds, that fetched data is considered fresh.
+   * After this time, data is considered expired and will be refetched when requested.
+   * @default time.days(7)
+   */
   cacheTime?: number;
-  disableDataCache?: boolean;
+  /**
+   * If `true`, the store's caching mechanisms will be fully disabled, meaning that the store will
+   * always refetch data on every call to `fetch()`, and the fetched data will not be stored unless
+   * a `setData` function is provided.
+   *
+   * Disable caching if you always want fresh data on refetch.
+   * @default false
+   */
+  disableCache?: boolean;
+  /**
+   * When `true`, the store actively fetches and refetches data as needed.
+   * When `false`, the store will not automatically fetch data until explicitly enabled.
+   * @default true
+   */
   enabled?: boolean;
+  /**
+   * Parameters to be passed to the fetcher, defined as either direct values or `ParamResolvable` functions.
+   * Dynamic parameters using `AttachValue` will cause the store to refetch when their values change.
+   */
   params?: {
     [K in keyof TParams]: ParamResolvable<TParams[K]>;
   };
+  /**
+   * The duration, in milliseconds, that data is considered fresh after fetching.
+   * After becoming stale, the store may automatically refetch data in the background if there are active subscribers.
+   *
+   * **Note:** Stale times under 5 seconds are strongly discouraged.
+   * @default time.minutes(2)
+   */
   staleTime?: number;
 };
 
 /**
- * A function that resolves to a value or an AttachValue wrapper.
+ * Represents a parameter that can be provided directly or defined via a reactive `AttachValue`.
+ * A parameter can be:
+ * - A static value (e.g. `string`, `number`).
+ * - A function that returns an `AttachValue<T>` when given a `SignalFunction`.
  */
 type ParamResolvable<T> = T | ((resolve: SignalFunction) => AttachValue<T>);
 
-/**
- * The result of resolving parameters into their direct values and AttachValue wrappers.
- */
 interface ResolvedParamsResult<TParams> {
+  /**
+   * Direct, non-reactive values resolved from the initial configuration.
+   */
   directValues: Partial<TParams>;
+  /**
+   * Reactive parameter values wrapped in `AttachValue`, which trigger refetches when they change.
+   */
   paramAttachVals: Partial<Record<keyof TParams, AttachValue<unknown>>>;
+  /**
+   * Fully resolved parameters, merging both direct and reactive values.
+   */
   resolvedParams: TParams;
 }
 
@@ -105,7 +225,6 @@ const [persist, discard] = [true, false];
 
 const SHOULD_PERSIST_INTERNAL_STATE_MAP: Record<string, boolean> = {
   /* Internal state to persist if the store is persisted */
-  data: persist,
   enabled: persist,
   error: persist,
   lastFetchedAt: persist,
@@ -114,6 +233,7 @@ const SHOULD_PERSIST_INTERNAL_STATE_MAP: Record<string, boolean> = {
 
   /* Internal state and methods to discard */
   fetch: discard,
+  getData: discard,
   isDataExpired: discard,
   isStale: discard,
   reset: discard,
@@ -137,7 +257,7 @@ const MIN_STALE_TIME = time.seconds(5);
  * @template U - User-defined custom store state
  * @template TData - The transformed data type, if applicable (defaults to `TQueryFnData`)
  */
-export function createRainbowQueryStore<
+export function createQueryStore<
   TQueryFnData,
   TParams extends Record<string, unknown> = Record<string, unknown>,
   U = unknown,
@@ -158,13 +278,15 @@ export function createRainbowQueryStore<
     onFetched,
     transform,
     cacheTime = time.days(7),
-    disableDataCache = true,
+    disableCache = false,
     enabled = true,
     params,
+    setData,
     staleTime = time.minutes(2),
+    suppressStaleTimeWarning = false,
   } = config;
 
-  if (IS_DEV && staleTime < MIN_STALE_TIME) {
+  if (IS_DEV && !suppressStaleTimeWarning && staleTime < MIN_STALE_TIME) {
     console.warn(
       `[RainbowQueryStore${persistConfig?.storageKey ? `: ${persistConfig.storageKey}` : ''}] ❌ Stale times under ${
         MIN_STALE_TIME / 1000
@@ -186,7 +308,6 @@ export function createRainbowQueryStore<
   let lastFetchKey: string | null = null;
 
   const initialData = {
-    data: null,
     enabled,
     error: null,
     lastFetchedAt: null,
@@ -207,32 +328,34 @@ export function createRainbowQueryStore<
     return currentParams as TParams;
   };
 
-  const scheduleNextFetch = (params: TParams) => {
-    if (staleTime <= 0) return;
-    if (activeRefetchTimeout) {
-      clearTimeout(activeRefetchTimeout);
-      activeRefetchTimeout = null;
-    }
-    activeRefetchTimeout = setTimeout(() => {
-      if (baseStore.getState().subscriptionCount > 0) {
-        baseStore.getState().fetch(params, { force: true });
-      }
-    }, staleTime);
-  };
-
   const createState: StateCreator<S, [], [['zustand/subscribeWithSelector', never]]> = (set, get, api) => {
     const pruneCache = (state: S): S => {
-      if (disableDataCache) return state;
       const now = Date.now();
       const newCache: Record<string, CacheEntry<TData>> = {};
       Object.entries(state.queryCache).forEach(([key, entry]) => {
-        if (now - entry.lastFetchedAt <= cacheTime) {
+        if (entry && now - entry.lastFetchedAt <= cacheTime) {
           newCache[key] = entry;
         }
       });
       return { ...state, queryCache: newCache };
     };
 
+    const scheduleNextFetch = (params: TParams) => {
+      if (staleTime <= 0) return;
+      if (activeRefetchTimeout) {
+        clearTimeout(activeRefetchTimeout);
+        activeRefetchTimeout = null;
+      }
+      const lastFetchedAt = get().queryCache[getQueryKey(params)]?.lastFetchedAt;
+      const timeUntilRefetch = lastFetchedAt ? staleTime - (Date.now() - lastFetchedAt) : staleTime;
+
+      activeRefetchTimeout = setTimeout(() => {
+        if (get().subscriptionCount > 0) {
+          get().fetch(params, { force: true });
+        }
+      }, timeUntilRefetch);
+    };
+
     const baseMethods = {
       async fetch(params: TParams | undefined, options: FetchOptions | undefined) {
         if (!get().enabled) return;
@@ -244,10 +367,9 @@ export function createRainbowQueryStore<
           return activeFetchPromise;
         }
 
-        if (!options?.force && !disableDataCache) {
-          const cached = get().queryCache[currentQueryKey];
-          if (cached && Date.now() - cached.lastFetchedAt <= (options?.staleTime ?? staleTime)) {
-            set(state => ({ ...state, data: cached.data }));
+        if (!options?.force && !disableCache) {
+          const lastFetchedAt = get().queryCache[currentQueryKey]?.lastFetchedAt;
+          if (lastFetchedAt && Date.now() - lastFetchedAt <= (options?.staleTime ?? staleTime)) {
             return;
           }
         }
@@ -262,32 +384,43 @@ export function createRainbowQueryStore<
             try {
               transformedData = transform ? transform(rawResult) : (rawResult as TData);
             } catch (transformError) {
-              throw new RainbowError(`[createRainbowQueryStore: ${persistConfig?.storageKey || currentQueryKey}]: transform failed`, {
+              throw new RainbowError(`[createQueryStore: ${persistConfig?.storageKey || currentQueryKey}]: transform failed`, {
                 cause: transformError,
               });
             }
 
             set(state => {
-              const newState = {
+              const lastFetchedAt = Date.now();
+              let newState: S = {
                 ...state,
                 error: null,
-                lastFetchedAt: Date.now(),
+                lastFetchedAt,
                 status: 'success' as const,
               };
 
-              if (!disableDataCache) {
+              if (!setData && !disableCache) {
                 newState.queryCache = {
                   ...newState.queryCache,
                   [currentQueryKey]: {
                     data: transformedData,
-                    lastFetchedAt: Date.now(),
+                    lastFetchedAt,
                   },
                 };
+              } else if (setData) {
+                setData(transformedData, (partial: S | Partial<S> | ((state: S) => S | Partial<S>)) => {
+                  newState = typeof partial === 'function' ? { ...newState, ...partial(newState) } : { ...newState, ...partial };
+                });
+                if (!disableCache) {
+                  newState.queryCache = {
+                    [currentQueryKey]: {
+                      data: null,
+                      lastFetchedAt,
+                    },
+                  };
+                }
               }
 
-              if (!onFetched) newState.data = transformedData;
-
-              return pruneCache(newState);
+              return disableCache ? newState : pruneCache(newState);
             });
 
             scheduleNextFetch(effectiveParams);
@@ -297,18 +430,16 @@ export function createRainbowQueryStore<
                 onFetched(transformedData, queryCapableStore);
               } catch (onFetchedError) {
                 logger.error(
-                  new RainbowError(
-                    `[createRainbowQueryStore: ${persistConfig?.storageKey || currentQueryKey}]: onFetched callback failed`,
-                    { cause: onFetchedError }
-                  )
+                  new RainbowError(`[createQueryStore: ${persistConfig?.storageKey || currentQueryKey}]: onFetched callback failed`, {
+                    cause: onFetchedError,
+                  })
                 );
               }
             }
           } catch (error) {
-            logger.error(
-              new RainbowError(`[createRainbowQueryStore: ${persistConfig?.storageKey || currentQueryKey}]: Failed to fetch data`),
-              { error }
-            );
+            logger.error(new RainbowError(`[createQueryStore: ${persistConfig?.storageKey || currentQueryKey}]: Failed to fetch data`), {
+              error,
+            });
             set(state => ({ ...state, error: error as Error, status: 'error' as const }));
             scheduleNextFetch(effectiveParams);
           } finally {
@@ -321,20 +452,26 @@ export function createRainbowQueryStore<
         return activeFetchPromise;
       },
 
-      isStale(staleTimeOverride?: number) {
-        const { lastFetchedAt } = get();
-        const effectiveStaleTime = staleTimeOverride ?? staleTime;
-        if (lastFetchedAt === null) return true;
-        return Date.now() - lastFetchedAt > effectiveStaleTime;
+      getData(params?: TParams) {
+        if (disableCache) return null;
+        const currentQueryKey = getQueryKey(params ?? getCurrentResolvedParams());
+        return get().queryCache[currentQueryKey]?.data ?? null;
       },
 
       isDataExpired(cacheTimeOverride?: number) {
-        const { lastFetchedAt } = get();
+        const lastFetchedAt = get().queryCache[getQueryKey(getCurrentResolvedParams())]?.lastFetchedAt;
         const effectiveCacheTime = cacheTimeOverride ?? cacheTime;
-        if (lastFetchedAt === null) return true;
+        if (!lastFetchedAt) return true;
         return Date.now() - lastFetchedAt > effectiveCacheTime;
       },
 
+      isStale(staleTimeOverride?: number) {
+        const lastFetchedAt = get().queryCache[getQueryKey(getCurrentResolvedParams())]?.lastFetchedAt;
+        const effectiveStaleTime = staleTimeOverride ?? staleTime;
+        if (!lastFetchedAt) return true;
+        return Date.now() - lastFetchedAt > effectiveStaleTime;
+      },
+
       reset() {
         if (activeRefetchTimeout) {
           clearTimeout(activeRefetchTimeout);
@@ -356,13 +493,14 @@ export function createRainbowQueryStore<
       const handleSetEnabled = subscribeWithSelector((state: S, prev: S) => {
         if (state.enabled !== prev.enabled) {
           if (state.enabled) {
-            const currentKey = getQueryKey(getCurrentResolvedParams());
+            const currentParams = getCurrentResolvedParams();
+            const currentKey = getQueryKey(currentParams);
             if (currentKey !== lastFetchKey) {
-              state.fetch(getCurrentResolvedParams(), { force: true });
-            } else if (!state.data || state.isStale()) {
+              state.fetch(currentParams, { force: true });
+            } else if (!state.queryCache[currentKey] || state.isStale()) {
               state.fetch();
             } else {
-              scheduleNextFetch(getCurrentResolvedParams());
+              scheduleNextFetch(currentParams);
             }
           } else {
             if (activeRefetchTimeout) {
@@ -373,9 +511,9 @@ export function createRainbowQueryStore<
         }
       });
 
-      const { data, fetch, isStale } = get();
+      const { fetch, isStale } = get();
 
-      if (!data || isStale()) {
+      if (!get().queryCache[getQueryKey(getCurrentResolvedParams())] || isStale()) {
         fetch(getCurrentResolvedParams(), { force: true });
       } else {
         scheduleNextFetch(getCurrentResolvedParams());
@@ -417,10 +555,6 @@ export function createRainbowQueryStore<
     : create<StoreState<TData, TParams> & U>()(subscribeWithSelector(createState));
 
   const queryCapableStore: QueryStore<TData, TParams, S> = Object.assign(baseStore, {
-    fetch: (params?: TParams, options?: FetchOptions) => baseStore.getState().fetch(params, options),
-    isDataExpired: (override?: number) => baseStore.getState().isDataExpired(override),
-    isStale: (override?: number) => baseStore.getState().isStale(override),
-    reset: () => baseStore.getState().reset(),
     enabled,
     destroy: () => {
       for (const unsub of paramUnsubscribes) {
@@ -429,6 +563,11 @@ export function createRainbowQueryStore<
       paramUnsubscribes.length = 0;
       queryCapableStore.getState().reset();
     },
+    fetch: (params?: TParams, options?: FetchOptions) => baseStore.getState().fetch(params, options),
+    getData: () => baseStore.getState().getData(),
+    isDataExpired: (override?: number) => baseStore.getState().isDataExpired(override),
+    isStale: (override?: number) => baseStore.getState().isStale(override),
+    reset: () => baseStore.getState().reset(),
   });
 
   const onParamChange = () => {
diff --git a/src/state/internal/tests/RainbowQueryStoreTest.tsx b/src/state/internal/tests/QueryStoreTest.tsx
similarity index 91%
rename from src/state/internal/tests/RainbowQueryStoreTest.tsx
rename to src/state/internal/tests/QueryStoreTest.tsx
index 43fba79f800..98267c9148a 100644
--- a/src/state/internal/tests/RainbowQueryStoreTest.tsx
+++ b/src/state/internal/tests/QueryStoreTest.tsx
@@ -6,7 +6,7 @@ import { Text } from '@/design-system';
 import { SupportedCurrencyKey } from '@/references';
 import { queryUserAssets } from '@/__swaps__/screens/Swap/resources/assets/userAssets';
 import { ParsedAssetsDictByChain } from '@/__swaps__/types/assets';
-import { createRainbowQueryStore, time } from '../createRainbowQueryStore';
+import { createQueryStore, time } from '../createQueryStore';
 import { createRainbowStore } from '../createRainbowStore';
 
 const ENABLE_LOGS = false;
@@ -45,25 +45,25 @@ type TestStore = {
 type QueryParams = { address: Address; currency: SupportedCurrencyKey };
 
 function logFetchInfo(params: QueryParams) {
-  console.log('[🔄 logFetchInfo 🔄] Current params:', JSON.stringify(Object.values(params), null, 2));
   const formattedTimeWithSeconds = new Date(Date.now()).toLocaleTimeString('en-US', {
     hour: '2-digit',
     minute: '2-digit',
     second: '2-digit',
   });
-  console.log('[🔄 Requesting Fetch 🔄] Last fetch attempt:', formattedTimeWithSeconds, '\nParams:', {
+  console.log('[🔄 UserAssetsTest - logFetchInfo 🔄]', '\nTime:', formattedTimeWithSeconds, '\nParams:', {
     address: params.address,
     currency: params.currency,
+    raw: JSON.stringify(Object.values(params), null, 2),
   });
 }
 
-export const userAssetsTestStore = createRainbowQueryStore<ParsedAssetsDictByChain, QueryParams, TestStore>(
+export const userAssetsTestStore = createQueryStore<ParsedAssetsDictByChain, QueryParams, TestStore>(
   {
     fetcher: ({ address, currency }) => {
       if (ENABLE_LOGS) logFetchInfo({ address, currency });
       return queryUserAssets({ address, currency });
     },
-    onFetched: (data, store) => store.setState({ userAssets: data }),
+    setData: (data, set) => set({ userAssets: data }),
 
     params: {
       address: $ => $(useAddressStore).address,
@@ -81,7 +81,11 @@ export const userAssetsTestStore = createRainbowQueryStore<ParsedAssetsDictByCha
         .reduce((max, asset) => Math.max(max, Number(asset.balance.display)), 0),
 
     setUserAssets: (data: ParsedAssetsDictByChain) => set({ userAssets: data }),
-  })
+  }),
+
+  {
+    storageKey: 'userAssetsQueryStoreTest',
+  }
 );
 
 export const UserAssetsTest = memo(function UserAssetsTest() {
@@ -110,7 +114,7 @@ export const UserAssetsTest = memo(function UserAssetsTest() {
         <View style={styles.buttonGroup}>
           <ButtonPressAnimation
             onPress={() => {
-              const currentAddress = useAddressStore.getState().nestedAddressTest.address;
+              const currentAddress = useAddressStore.getState().address;
               switch (currentAddress) {
                 case testAddresses[0]:
                   useAddressStore.getState().setAddress(testAddresses[1]);

From a8d4c73c2f3d0baa3d48d8180e7813050d396f61 Mon Sep 17 00:00:00 2001
From: Christian Baroni <7061887+christianbaroni@users.noreply.github.com>
Date: Tue, 17 Dec 2024 21:52:48 +0000
Subject: [PATCH 13/29] Use consistent store naming

---
 src/state/internal/tests/QueryStoreTest.tsx | 12 ++++++------
 1 file changed, 6 insertions(+), 6 deletions(-)

diff --git a/src/state/internal/tests/QueryStoreTest.tsx b/src/state/internal/tests/QueryStoreTest.tsx
index 98267c9148a..736ca446a31 100644
--- a/src/state/internal/tests/QueryStoreTest.tsx
+++ b/src/state/internal/tests/QueryStoreTest.tsx
@@ -57,7 +57,7 @@ function logFetchInfo(params: QueryParams) {
   });
 }
 
-export const userAssetsTestStore = createQueryStore<ParsedAssetsDictByChain, QueryParams, TestStore>(
+export const useUserAssetsTestStore = createQueryStore<ParsedAssetsDictByChain, QueryParams, TestStore>(
   {
     fetcher: ({ address, currency }) => {
       if (ENABLE_LOGS) logFetchInfo({ address, currency });
@@ -89,8 +89,8 @@ export const userAssetsTestStore = createQueryStore<ParsedAssetsDictByChain, Que
 );
 
 export const UserAssetsTest = memo(function UserAssetsTest() {
-  const data = userAssetsTestStore(state => state.userAssets);
-  const enabled = userAssetsTestStore(state => state.enabled);
+  const data = useUserAssetsTestStore(state => state.userAssets);
+  const enabled = useUserAssetsTestStore(state => state.enabled);
 
   useEffect(() => {
     if (ENABLE_LOGS) {
@@ -135,12 +135,12 @@ export const UserAssetsTest = memo(function UserAssetsTest() {
           </ButtonPressAnimation>
           <ButtonPressAnimation
             onPress={() => {
-              userAssetsTestStore.setState({ enabled: !enabled });
+              useUserAssetsTestStore.setState({ enabled: !enabled });
             }}
             style={styles.button}
           >
             <Text color="label" size="17pt" weight="heavy">
-              {userAssetsTestStore.getState().enabled ? 'Disable Fetching' : 'Enable Fetching'}
+              {useUserAssetsTestStore.getState().enabled ? 'Disable Fetching' : 'Enable Fetching'}
             </Text>
           </ButtonPressAnimation>
         </View>
@@ -149,7 +149,7 @@ export const UserAssetsTest = memo(function UserAssetsTest() {
   );
 });
 
-if (ENABLE_LOGS) console.log('[💾 UserAssetsTest 💾] Initial data exists:', !!userAssetsTestStore.getState().userAssets);
+if (ENABLE_LOGS) console.log('[💾 UserAssetsTest 💾] Initial data exists:', !!useUserAssetsTestStore.getState().userAssets);
 
 const styles = StyleSheet.create({
   button: {

From 7dbb6ef3be11dccc36ef29ed37c014d55a8fe6e7 Mon Sep 17 00:00:00 2001
From: Christian Baroni <7061887+christianbaroni@users.noreply.github.com>
Date: Thu, 19 Dec 2024 17:58:25 +0000
Subject: [PATCH 14/29] Improve error handling, add getStatus(), reduce type
 redundancy, omit store methods from persisted state

---
 src/state/internal/createQueryStore.ts      | 226 +++++++++++++++-----
 src/state/internal/createRainbowStore.ts    |  21 +-
 src/state/internal/tests/QueryStoreTest.tsx |   7 +-
 3 files changed, 199 insertions(+), 55 deletions(-)

diff --git a/src/state/internal/createQueryStore.ts b/src/state/internal/createQueryStore.ts
index 4299e2a8fda..a63b06db04f 100644
--- a/src/state/internal/createQueryStore.ts
+++ b/src/state/internal/createQueryStore.ts
@@ -1,9 +1,9 @@
 import { StateCreator, StoreApi, UseBoundStore, create } from 'zustand';
 import { subscribeWithSelector } from 'zustand/middleware';
 import { IS_DEV } from '@/env';
-import { logger, RainbowError } from '@/logger';
-import { createRainbowStore, RainbowPersistConfig } from './createRainbowStore';
-import { $, AttachValue, attachValueSubscriptionMap, SignalFunction, Unsubscribe } from './signal';
+import { RainbowError, logger } from '@/logger';
+import { RainbowPersistConfig, createRainbowStore, omitStoreMethods } from './createRainbowStore';
+import { $, AttachValue, SignalFunction, Unsubscribe, attachValueSubscriptionMap } from './signal';
 
 const ENABLE_LOGS = false;
 
@@ -28,14 +28,25 @@ export const QueryStatuses = {
  */
 export type QueryStatus = (typeof QueryStatuses)[keyof typeof QueryStatuses];
 
+/**
+ * Expanded status information for the currently specified query parameters.
+ */
+export type QueryStatusInfo = {
+  isError: boolean;
+  isFetching: boolean;
+  isIdle: boolean;
+  isInitialLoading: boolean;
+  isSuccess: boolean;
+};
+
 /**
  * Defines additional options for a data fetch operation.
  */
 interface FetchOptions {
   /**
    * Overrides the default cache duration for this fetch, in milliseconds.
-   * When data in the cache is older than this duration, it will be considered expired and
-   * become eligible for pruning.
+   * If data in the cache is older than this duration, it will be considered expired and
+   * will be pruned following a successful fetch.
    */
   cacheTime?: number;
   /**
@@ -45,17 +56,23 @@ interface FetchOptions {
   force?: boolean;
   /**
    * Overrides the default stale duration for this fetch, in milliseconds.
-   * When data is older than this duration, it is considered stale and if the query is active,
-   * a background refetch will occur.
+   * If the fetch is successful, the subsequently scheduled refetch will occur after
+   * the specified duration.
    */
   staleTime?: number;
 }
 
 /**
- * Represents an entry in the query cache, which stores fetched data along with metadata.
+ * Represents an entry in the query cache, which stores fetched data along with metadata, and error information
+ * in the event the most recent fetch failed.
  */
 interface CacheEntry<TData> {
   data: TData | null;
+  errorInfo: {
+    error: Error;
+    lastFailedAt: number;
+    retryCount: number;
+  } | null;
   lastFetchedAt: number;
 }
 
@@ -65,6 +82,8 @@ interface CacheEntry<TData> {
  * In addition to Zustand's store API (such as `getState()` and `subscribe()`), this interface provides:
  * - **`enabled`**: A boolean indicating if the store is actively fetching data.
  * - **`fetch(params, options)`**: Initiates a data fetch operation.
+ * - **`getData(params)`**: Returns the cached data, if available, for the current query parameters.
+ * - **`getStatus()`**: Returns expanded status information for the current query parameters.
  * - **`isDataExpired(override?)`**: Checks if the current data has expired based on `cacheTime`.
  * - **`isStale(override?)`**: Checks if the current data is stale based on `staleTime`.
  * - **`reset()`**: Resets the store to its initial state, clearing data and errors.
@@ -89,6 +108,16 @@ export interface QueryStore<TData, TParams extends Record<string, unknown>, S ex
    * @returns The cached data, or `null` if no data is available.
    */
   getData: (params?: TParams) => TData | null;
+  /**
+   * Returns expanded status information for the currently specified query parameters. The raw
+   * status can be obtained by directly reading the `status` property.
+   * @example
+   * ```ts
+   * const isInitialLoad = useQueryStore(state => state.getStatus().isInitialLoad);
+   * ```
+   * @returns An object containing boolean flags for each status.
+   */
+  getStatus: () => QueryStatusInfo;
   /**
    * Determines if the current data is expired, meaning it has exceeded the `cacheTime` duration.
    * @param override - An optional override for the default cache time, in milliseconds.
@@ -112,18 +141,15 @@ export interface QueryStore<TData, TParams extends Record<string, unknown>, S ex
  * The state structure managed by the query store, including query-related fields and actions.
  * This type is generally internal, and extended by user-defined states when creating a store.
  */
-type StoreState<TData, TParams extends Record<string, unknown>> = {
-  enabled: boolean;
+type StoreState<TData, TParams extends Record<string, unknown>> = Pick<
+  QueryStore<TData, TParams, StoreState<TData, TParams>>,
+  'enabled' | 'fetch' | 'getData' | 'getStatus' | 'isDataExpired' | 'isStale' | 'reset'
+> & {
   error: Error | null;
   lastFetchedAt: number | null;
   queryCache: Record<string, CacheEntry<TData> | undefined>;
   status: QueryStatus;
   subscriptionCount: number;
-  fetch: (params?: TParams, options?: FetchOptions) => Promise<void>;
-  getData: (params?: TParams) => TData | null;
-  isDataExpired: (cacheTimeOverride?: number) => boolean;
-  isStale: (staleTimeOverride?: number) => boolean;
-  reset: () => void;
 };
 
 /**
@@ -135,6 +161,22 @@ export type RainbowQueryStoreConfig<TQueryFnData, TParams extends Record<string,
    * Receives parameters of type `TParams` and returns either a promise or a raw data value of type `TQueryFnData`.
    */
   fetcher: (params: TParams) => TQueryFnData | Promise<TQueryFnData>;
+  /**
+   * The maximum number of times to retry a failed fetch operation.
+   * @default 3
+   */
+  maxRetries?: number;
+  /**
+   * The delay between retries after a fetch error occurs, in milliseconds, defined as a number or a function that
+   * receives the error and current retry count and returns a number.
+   * @default time.seconds(5)
+   */
+  retryDelay?: number | ((retryCount: number, error: Error) => number);
+  /**
+   * A callback invoked whenever a fetch operation fails.
+   * Receives the error and the current retry count.
+   */
+  onError?: (error: Error, retryCount: number) => void;
   /**
    * A callback invoked whenever fresh data is successfully fetched.
    * Receives the transformed data and the store instance, allowing for side effects or additional updates.
@@ -234,6 +276,7 @@ const SHOULD_PERSIST_INTERNAL_STATE_MAP: Record<string, boolean> = {
   /* Internal state and methods to discard */
   fetch: discard,
   getData: discard,
+  getStatus: discard,
   isDataExpired: discard,
   isStale: discard,
   reset: discard,
@@ -280,7 +323,10 @@ export function createQueryStore<
     cacheTime = time.days(7),
     disableCache = false,
     enabled = true,
+    maxRetries = 3,
+    onError,
     params,
+    retryDelay = time.seconds(5),
     setData,
     staleTime = time.minutes(2),
     suppressStaleTimeWarning = false,
@@ -312,7 +358,7 @@ export function createQueryStore<
     error: null,
     lastFetchedAt: null,
     queryCache: {},
-    status: 'idle' as const,
+    status: QueryStatuses.Idle,
     subscriptionCount: 0,
   };
 
@@ -340,14 +386,17 @@ export function createQueryStore<
       return { ...state, queryCache: newCache };
     };
 
-    const scheduleNextFetch = (params: TParams) => {
-      if (staleTime <= 0) return;
+    const scheduleNextFetch = (params: TParams, options: FetchOptions | undefined) => {
+      const effectiveStaleTime = options?.staleTime ?? staleTime;
+      if (effectiveStaleTime <= 0 || effectiveStaleTime === Infinity) return;
       if (activeRefetchTimeout) {
         clearTimeout(activeRefetchTimeout);
         activeRefetchTimeout = null;
       }
-      const lastFetchedAt = get().queryCache[getQueryKey(params)]?.lastFetchedAt;
-      const timeUntilRefetch = lastFetchedAt ? staleTime - (Date.now() - lastFetchedAt) : staleTime;
+      const currentQueryKey = getQueryKey(params);
+      const lastFetchedAt =
+        get().queryCache[currentQueryKey]?.lastFetchedAt || (disableCache && lastFetchKey === currentQueryKey ? get().lastFetchedAt : null);
+      const timeUntilRefetch = lastFetchedAt ? effectiveStaleTime - (Date.now() - lastFetchedAt) : effectiveStaleTime;
 
       activeRefetchTimeout = setTimeout(() => {
         if (get().subscriptionCount > 0) {
@@ -361,20 +410,21 @@ export function createQueryStore<
         if (!get().enabled) return;
         const effectiveParams = params ?? getCurrentResolvedParams();
         const currentQueryKey = getQueryKey(effectiveParams);
-        const isLoading = get().status === 'loading';
+        const isLoading = get().status === QueryStatuses.Loading;
 
         if (activeFetchPromise && lastFetchKey === currentQueryKey && isLoading && !options?.force) {
           return activeFetchPromise;
         }
 
         if (!options?.force && !disableCache) {
-          const lastFetchedAt = get().queryCache[currentQueryKey]?.lastFetchedAt;
-          if (lastFetchedAt && Date.now() - lastFetchedAt <= (options?.staleTime ?? staleTime)) {
+          const { errorInfo, lastFetchedAt } = get().queryCache[currentQueryKey] ?? {};
+          const errorRetriesExhausted = errorInfo && errorInfo.retryCount >= maxRetries;
+          if ((!errorInfo || errorRetriesExhausted) && lastFetchedAt && Date.now() - lastFetchedAt <= (options?.staleTime ?? staleTime)) {
             return;
           }
         }
 
-        set(state => ({ ...state, error: null, status: 'loading' }));
+        set(state => ({ ...state, error: null, status: QueryStatuses.Loading }));
         lastFetchKey = currentQueryKey;
 
         const fetchOperation = async () => {
@@ -395,7 +445,7 @@ export function createQueryStore<
                 ...state,
                 error: null,
                 lastFetchedAt,
-                status: 'success' as const,
+                status: QueryStatuses.Success,
               };
 
               if (!setData && !disableCache) {
@@ -403,6 +453,7 @@ export function createQueryStore<
                   ...newState.queryCache,
                   [currentQueryKey]: {
                     data: transformedData,
+                    errorInfo: null,
                     lastFetchedAt,
                   },
                 };
@@ -414,16 +465,17 @@ export function createQueryStore<
                   newState.queryCache = {
                     [currentQueryKey]: {
                       data: null,
+                      errorInfo: null,
                       lastFetchedAt,
                     },
                   };
                 }
               }
 
-              return disableCache ? newState : pruneCache(newState);
+              return disableCache || cacheTime === Infinity ? newState : pruneCache(newState);
             });
 
-            scheduleNextFetch(effectiveParams);
+            scheduleNextFetch(effectiveParams, options);
 
             if (onFetched) {
               try {
@@ -437,11 +489,58 @@ export function createQueryStore<
               }
             }
           } catch (error) {
+            const typedError = error instanceof Error ? error : new Error(String(error));
+            const entry = get().queryCache[currentQueryKey];
+            const currentRetryCount = entry?.errorInfo?.retryCount ?? 0;
+
+            onError?.(typedError, currentRetryCount);
+
+            if (currentRetryCount < maxRetries) {
+              const errorRetryDelay = typeof retryDelay === 'function' ? retryDelay(currentRetryCount, typedError) : retryDelay;
+
+              if (get().subscriptionCount > 0) {
+                activeRefetchTimeout = setTimeout(() => {
+                  if (get().subscriptionCount > 0) {
+                    get().fetch(effectiveParams, { force: true });
+                  }
+                }, errorRetryDelay);
+              }
+
+              set(state => ({
+                ...state,
+                error: typedError,
+                status: QueryStatuses.Error,
+                queryCache: {
+                  ...state.queryCache,
+                  [currentQueryKey]: {
+                    ...entry,
+                    errorState: {
+                      error: typedError,
+                      retryCount: currentRetryCount + 1,
+                    },
+                  },
+                },
+              }));
+            } else {
+              set(state => ({
+                ...state,
+                status: QueryStatuses.Error,
+                queryCache: {
+                  ...state.queryCache,
+                  [currentQueryKey]: {
+                    ...entry,
+                    errorState: {
+                      error: typedError,
+                      retryCount: currentRetryCount,
+                    },
+                  },
+                },
+              }));
+            }
+
             logger.error(new RainbowError(`[createQueryStore: ${persistConfig?.storageKey || currentQueryKey}]: Failed to fetch data`), {
-              error,
+              error: typedError,
             });
-            set(state => ({ ...state, error: error as Error, status: 'error' as const }));
-            scheduleNextFetch(effectiveParams);
           } finally {
             activeFetchPromise = null;
             lastFetchKey = null;
@@ -458,17 +557,41 @@ export function createQueryStore<
         return get().queryCache[currentQueryKey]?.data ?? null;
       },
 
+      getStatus() {
+        const status = get().status;
+        const currentQueryKey = getQueryKey(getCurrentResolvedParams());
+        const lastFetchedAt =
+          get().queryCache[currentQueryKey]?.lastFetchedAt ||
+          (disableCache && lastFetchKey === currentQueryKey ? get().lastFetchedAt : null);
+
+        return {
+          isError: status === QueryStatuses.Error,
+          isFetching: status === QueryStatuses.Loading,
+          isIdle: status === QueryStatuses.Idle,
+          isInitialLoading: !lastFetchedAt && status === QueryStatuses.Loading,
+          isSuccess: status === QueryStatuses.Success,
+        };
+      },
+
       isDataExpired(cacheTimeOverride?: number) {
-        const lastFetchedAt = get().queryCache[getQueryKey(getCurrentResolvedParams())]?.lastFetchedAt;
-        const effectiveCacheTime = cacheTimeOverride ?? cacheTime;
+        const currentQueryKey = getQueryKey(getCurrentResolvedParams());
+        const lastFetchedAt =
+          get().queryCache[currentQueryKey]?.lastFetchedAt ||
+          (disableCache && lastFetchKey === currentQueryKey ? get().lastFetchedAt : null);
+
         if (!lastFetchedAt) return true;
+        const effectiveCacheTime = cacheTimeOverride ?? cacheTime;
         return Date.now() - lastFetchedAt > effectiveCacheTime;
       },
 
       isStale(staleTimeOverride?: number) {
-        const lastFetchedAt = get().queryCache[getQueryKey(getCurrentResolvedParams())]?.lastFetchedAt;
-        const effectiveStaleTime = staleTimeOverride ?? staleTime;
+        const currentQueryKey = getQueryKey(getCurrentResolvedParams());
+        const lastFetchedAt =
+          get().queryCache[currentQueryKey]?.lastFetchedAt ||
+          (disableCache && lastFetchKey === currentQueryKey ? get().lastFetchedAt : null);
+
         if (!lastFetchedAt) return true;
+        const effectiveStaleTime = staleTimeOverride ?? staleTime;
         return Date.now() - lastFetchedAt > effectiveStaleTime;
       },
 
@@ -483,15 +606,14 @@ export function createQueryStore<
       },
     };
 
-    const userState = customStateCreator?.(set, get, api) ?? ({} as U);
-
     const subscribeWithSelector = api.subscribe;
+
     api.subscribe = (listener: (state: S, prevState: S) => void) => {
-      set(prev => ({ ...prev, subscriptionCount: prev.subscriptionCount + 1 }));
+      set(state => ({ ...state, subscriptionCount: state.subscriptionCount + 1 }));
       const unsubscribe = subscribeWithSelector(listener);
 
-      const handleSetEnabled = subscribeWithSelector((state: S, prev: S) => {
-        if (state.enabled !== prev.enabled) {
+      const handleSetEnabled = subscribeWithSelector((state: S, prevState: S) => {
+        if (state.enabled !== prevState.enabled) {
           if (state.enabled) {
             const currentParams = getCurrentResolvedParams();
             const currentKey = getQueryKey(currentParams);
@@ -500,7 +622,7 @@ export function createQueryStore<
             } else if (!state.queryCache[currentKey] || state.isStale()) {
               state.fetch();
             } else {
-              scheduleNextFetch(currentParams);
+              scheduleNextFetch(currentParams, undefined);
             }
           } else {
             if (activeRefetchTimeout) {
@@ -512,29 +634,32 @@ export function createQueryStore<
       });
 
       const { fetch, isStale } = get();
+      const currentParams = getCurrentResolvedParams()
 
-      if (!get().queryCache[getQueryKey(getCurrentResolvedParams())] || isStale()) {
-        fetch(getCurrentResolvedParams(), { force: true });
+      if (!get().queryCache[getQueryKey(currentParams)] || isStale()) {
+        fetch(currentParams, { force: true });
       } else {
-        scheduleNextFetch(getCurrentResolvedParams());
+        scheduleNextFetch(currentParams, undefined);
       }
 
       return () => {
         handleSetEnabled();
         unsubscribe();
-        set(prev => {
-          const newCount = Math.max(prev.subscriptionCount - 1, 0);
+        set(state => {
+          const newCount = Math.max(state.subscriptionCount - 1, 0);
           if (newCount === 0) {
             if (activeRefetchTimeout) {
               clearTimeout(activeRefetchTimeout);
               activeRefetchTimeout = null;
             }
           }
-          return { ...prev, subscriptionCount: newCount };
+          return { ...state, subscriptionCount: newCount };
         });
       };
     };
 
+    const userState = customStateCreator?.(set, get, api) ?? ({} as U);
+
     /* Merge base data, user state, and methods into the final store state */
     return {
       ...initialData,
@@ -565,6 +690,7 @@ export function createQueryStore<
     },
     fetch: (params?: TParams, options?: FetchOptions) => baseStore.getState().fetch(params, options),
     getData: () => baseStore.getState().getData(),
+    getStatus: () => baseStore.getState().getStatus(),
     isDataExpired: (override?: number) => baseStore.getState().isDataExpired(override),
     isStale: (override?: number) => baseStore.getState().isStale(override),
     reset: () => baseStore.getState().reset(),
@@ -608,9 +734,9 @@ function isParamFn<T>(param: T | ((resolve: SignalFunction) => AttachValue<T>)):
 function resolveParams<TParams extends Record<string, unknown>>(params: {
   [K in keyof TParams]: ParamResolvable<TParams[K]>;
 }): ResolvedParamsResult<TParams> {
-  const resolvedParams = {} as TParams;
-  const paramAttachVals: Partial<Record<keyof TParams, AttachValue<unknown>>> = {};
   const directValues: Partial<TParams> = {};
+  const paramAttachVals: Partial<Record<keyof TParams, AttachValue<unknown>>> = {};
+  const resolvedParams = {} as TParams;
 
   for (const key in params) {
     const param = params[key];
@@ -624,7 +750,7 @@ function resolveParams<TParams extends Record<string, unknown>>(params: {
     }
   }
 
-  return { resolvedParams, paramAttachVals, directValues };
+  return { directValues, paramAttachVals, resolvedParams };
 }
 
 function createBlendedPartialize<TData, TParams extends Record<string, unknown>, S extends StoreState<TData, TParams> & U, U = unknown>(
@@ -642,7 +768,7 @@ function createBlendedPartialize<TData, TParams extends Record<string, unknown>,
     }
 
     return {
-      ...(userPartialize ? userPartialize(clonedState) : clonedState),
+      ...(userPartialize ? userPartialize(clonedState) : omitStoreMethods(clonedState)),
       ...internalStateToPersist,
     } satisfies Partial<S>;
   };
diff --git a/src/state/internal/createRainbowStore.ts b/src/state/internal/createRainbowStore.ts
index 5de0748c95d..edb2ee3ae0b 100644
--- a/src/state/internal/createRainbowStore.ts
+++ b/src/state/internal/createRainbowStore.ts
@@ -65,7 +65,7 @@ export function createRainbowStore<S = unknown>(
       persist(createState, {
         migrate: persistConfig.migrate,
         name: persistConfig.storageKey,
-        partialize: persistConfig.partialize || (state => state),
+        partialize: persistConfig.partialize || omitStoreMethods,
         storage: persistStorage,
         version,
       })
@@ -73,6 +73,23 @@ export function createRainbowStore<S = unknown>(
   );
 }
 
+/**
+ * Default partialize function if none is provided. It omits top-level store
+ * methods and keeps all other state.
+ */
+export function omitStoreMethods<S>(state: S): Partial<S> {
+  if (state !== null && typeof state === 'object') {
+    const result: Record<string, unknown> = {};
+    Object.entries(state).forEach(([key, val]) => {
+      if (typeof val !== 'function') {
+        result[key] = val;
+      }
+    });
+    return result as Partial<S>;
+  }
+  return state;
+}
+
 /**
  * Creates a persist storage object for the Rainbow store.
  * @param config - The configuration options for the persistable Rainbow store.
@@ -88,7 +105,7 @@ function createPersistStorage<S = unknown>(config: RainbowPersistConfig<S>) {
       if (!serializedValue) return null;
       return deserializer(serializedValue);
     },
-    setItem: (name, value) =>
+    setItem: (name: string, value: StorageValue<Partial<S>>) =>
       lazyPersist({
         serializer,
         storageKey,
diff --git a/src/state/internal/tests/QueryStoreTest.tsx b/src/state/internal/tests/QueryStoreTest.tsx
index 736ca446a31..14f3c39a708 100644
--- a/src/state/internal/tests/QueryStoreTest.tsx
+++ b/src/state/internal/tests/QueryStoreTest.tsx
@@ -1,3 +1,6 @@
+// ⚠️ Uncomment everything below to experiment with the QueryStore creator
+// TODO: Comment out test code below before merging
+
 import React, { memo, useEffect } from 'react';
 import { StyleSheet, View } from 'react-native';
 import { Address } from 'viem';
@@ -83,9 +86,7 @@ export const useUserAssetsTestStore = createQueryStore<ParsedAssetsDictByChain,
     setUserAssets: (data: ParsedAssetsDictByChain) => set({ userAssets: data }),
   }),
 
-  {
-    storageKey: 'userAssetsQueryStoreTest',
-  }
+  { storageKey: 'userAssetsQueryStoreTest' }
 );
 
 export const UserAssetsTest = memo(function UserAssetsTest() {

From 2f86899558e1527e5431a8ed660461667174db11 Mon Sep 17 00:00:00 2001
From: Christian Baroni <7061887+christianbaroni@users.noreply.github.com>
Date: Thu, 19 Dec 2024 17:59:51 +0000
Subject: [PATCH 15/29] [createRainbowStore] Support maps and sets internally

---
 src/state/internal/createRainbowStore.ts | 46 ++++++++++++++++++++++--
 1 file changed, 44 insertions(+), 2 deletions(-)

diff --git a/src/state/internal/createRainbowStore.ts b/src/state/internal/createRainbowStore.ts
index edb2ee3ae0b..195d167e8e2 100644
--- a/src/state/internal/createRainbowStore.ts
+++ b/src/state/internal/createRainbowStore.ts
@@ -157,7 +157,7 @@ const lazyPersist = <S>({ name, serializer, storageKey, value }: LazyPersistPara
  */
 function defaultSerializeState<S>(state: StorageValue<Partial<S>>['state'], version: StorageValue<Partial<S>>['version']): string {
   try {
-    return JSON.stringify({ state, version });
+    return JSON.stringify({ state, version }, replacer);
   } catch (error) {
     logger.error(new RainbowError(`[createRainbowStore]: Failed to serialize Rainbow store data`), { error });
     throw error;
@@ -171,9 +171,51 @@ function defaultSerializeState<S>(state: StorageValue<Partial<S>>['state'], vers
  */
 function defaultDeserializeState<S>(serializedState: string): StorageValue<Partial<S>> {
   try {
-    return JSON.parse(serializedState);
+    return JSON.parse(serializedState, reviver);
   } catch (error) {
     logger.error(new RainbowError(`[createRainbowStore]: Failed to deserialize persisted Rainbow store data`), { error });
     throw error;
   }
 }
+
+interface SerializedMap {
+  __type: 'Map';
+  entries: [unknown, unknown][];
+}
+
+function isSerializedMap(value: unknown): value is SerializedMap {
+  return typeof value === 'object' && value !== null && (value as Record<string, unknown>).__type === 'Map';
+}
+
+interface SerializedSet {
+  __type: 'Set';
+  values: unknown[];
+}
+
+function isSerializedSet(value: unknown): value is SerializedSet {
+  return typeof value === 'object' && value !== null && (value as Record<string, unknown>).__type === 'Set';
+}
+
+/**
+ * Replacer function to handle serialization of Maps and Sets.
+ */
+function replacer(key: string, value: unknown): unknown {
+  if (value instanceof Map) {
+    return { __type: 'Map', entries: Array.from(value.entries()) };
+  } else if (value instanceof Set) {
+    return { __type: 'Set', values: Array.from(value) };
+  }
+  return value;
+}
+
+/**
+ * Reviver function to handle deserialization of Maps and Sets.
+ */
+function reviver(key: string, value: unknown): unknown {
+  if (isSerializedMap(value)) {
+    return new Map(value.entries);
+  } else if (isSerializedSet(value)) {
+    return new Set(value.values);
+  }
+  return value;
+}

From 791757e9095f156ef5efbe0f6c5e620bf9fd2231 Mon Sep 17 00:00:00 2001
From: Christian Baroni <7061887+christianbaroni@users.noreply.github.com>
Date: Fri, 20 Dec 2024 04:55:44 +0000
Subject: [PATCH 16/29] Allow dynamic params for internal state: ($, store) =>
 $(store).state

---
 src/state/internal/createQueryStore.ts      | 55 ++++++++++-----------
 src/state/internal/signal.ts                | 20 ++++----
 src/state/internal/tests/QueryStoreTest.tsx | 45 ++++++++---------
 3 files changed, 55 insertions(+), 65 deletions(-)

diff --git a/src/state/internal/createQueryStore.ts b/src/state/internal/createQueryStore.ts
index a63b06db04f..463a1970a5b 100644
--- a/src/state/internal/createQueryStore.ts
+++ b/src/state/internal/createQueryStore.ts
@@ -179,12 +179,12 @@ export type RainbowQueryStoreConfig<TQueryFnData, TParams extends Record<string,
   onError?: (error: Error, retryCount: number) => void;
   /**
    * A callback invoked whenever fresh data is successfully fetched.
-   * Receives the transformed data and the store instance, allowing for side effects or additional updates.
+   * Receives the transformed data and the store's set function, which can optionally be used to update store state.
    */
-  onFetched?: (data: TData, store: QueryStore<TData, TParams, S>) => void;
+  onFetched?: (data: TData, set: (partial: S | Partial<S> | ((state: S) => S | Partial<S>)) => void) => void;
   /**
    * A function that overrides the default behavior of setting the fetched data in the store's query cache.
-   * Receives the transformed data and a set function that updates the store state.
+   * Receives the transformed data and the store's set function.
    */
   setData?: (data: TData, set: (partial: S | Partial<S> | ((state: S) => S | Partial<S>)) => void) => void;
   /**
@@ -223,7 +223,7 @@ export type RainbowQueryStoreConfig<TQueryFnData, TParams extends Record<string,
    * Dynamic parameters using `AttachValue` will cause the store to refetch when their values change.
    */
   params?: {
-    [K in keyof TParams]: ParamResolvable<TParams[K]>;
+    [K in keyof TParams]: ParamResolvable<TParams[K], TParams, S, TData>;
   };
   /**
    * The duration, in milliseconds, that data is considered fresh after fetching.
@@ -241,7 +241,9 @@ export type RainbowQueryStoreConfig<TQueryFnData, TParams extends Record<string,
  * - A static value (e.g. `string`, `number`).
  * - A function that returns an `AttachValue<T>` when given a `SignalFunction`.
  */
-type ParamResolvable<T> = T | ((resolve: SignalFunction) => AttachValue<T>);
+type ParamResolvable<T, TParams extends Record<string, unknown>, S extends StoreState<TData, TParams>, TData> =
+  | T
+  | (($: SignalFunction, store: QueryStore<TData, TParams, S>) => AttachValue<T>);
 
 interface ResolvedParamsResult<TParams> {
   /**
@@ -307,9 +309,7 @@ export function createQueryStore<
   TData = TQueryFnData,
 >(
   config: RainbowQueryStoreConfig<TQueryFnData, TParams, TData, StoreState<TData, TParams> & U> & {
-    params?: {
-      [K in keyof TParams]: ParamResolvable<TParams[K]>;
-    };
+    params?: { [K in keyof TParams]: ParamResolvable<TParams[K], TParams, StoreState<TData, TParams> & U, TData> };
   },
   customStateCreator?: StateCreator<U, [], [['zustand/subscribeWithSelector', never]]>,
   persistConfig?: RainbowPersistConfig<StoreState<TData, TParams> & U>
@@ -343,12 +343,6 @@ export function createQueryStore<
   let directValues: Partial<TParams> = {};
   let paramAttachVals: Partial<Record<keyof TParams, AttachValue<unknown>>> = {};
 
-  if (params) {
-    const result = resolveParams<TParams>(params);
-    paramAttachVals = result.paramAttachVals;
-    directValues = result.directValues;
-  }
-
   let activeFetchPromise: Promise<void> | null = null;
   let activeRefetchTimeout: NodeJS.Timeout | null = null;
   let lastFetchKey: string | null = null;
@@ -479,7 +473,7 @@ export function createQueryStore<
 
             if (onFetched) {
               try {
-                onFetched(transformedData, queryCapableStore);
+                onFetched(transformedData, set);
               } catch (onFetchedError) {
                 logger.error(
                   new RainbowError(`[createQueryStore: ${persistConfig?.storageKey || currentQueryKey}]: onFetched callback failed`, {
@@ -634,7 +628,7 @@ export function createQueryStore<
       });
 
       const { fetch, isStale } = get();
-      const currentParams = getCurrentResolvedParams()
+      const currentParams = getCurrentResolvedParams();
 
       if (!get().queryCache[getQueryKey(currentParams)] || isStale()) {
         fetch(currentParams, { force: true });
@@ -682,10 +676,8 @@ export function createQueryStore<
   const queryCapableStore: QueryStore<TData, TParams, S> = Object.assign(baseStore, {
     enabled,
     destroy: () => {
-      for (const unsub of paramUnsubscribes) {
-        unsub();
-      }
-      paramUnsubscribes.length = 0;
+      for (const unsub of paramUnsubscribes) unsub();
+      paramUnsubscribes = [];
       queryCapableStore.getState().reset();
     },
     fetch: (params?: TParams, options?: FetchOptions) => baseStore.getState().fetch(params, options),
@@ -696,12 +688,18 @@ export function createQueryStore<
     reset: () => baseStore.getState().reset(),
   });
 
+  if (params) {
+    const result = resolveParams<TParams, S, TData>(params, queryCapableStore);
+    paramAttachVals = result.paramAttachVals;
+    directValues = result.directValues;
+  }
+
   const onParamChange = () => {
     const newParams = getCurrentResolvedParams();
     queryCapableStore.fetch(newParams, { force: true });
   };
 
-  const paramUnsubscribes: Unsubscribe[] = [];
+  let paramUnsubscribes: Unsubscribe[] = [];
 
   for (const k in paramAttachVals) {
     const attachVal = paramAttachVals[k];
@@ -727,21 +725,18 @@ export function createQueryStore<
   return queryCapableStore;
 }
 
-function isParamFn<T>(param: T | ((resolve: SignalFunction) => AttachValue<T>)): param is (resolve: SignalFunction) => AttachValue<T> {
-  return typeof param === 'function';
-}
-
-function resolveParams<TParams extends Record<string, unknown>>(params: {
-  [K in keyof TParams]: ParamResolvable<TParams[K]>;
-}): ResolvedParamsResult<TParams> {
+function resolveParams<TParams extends Record<string, unknown>, S extends StoreState<TData, TParams> & U, TData, U = unknown>(
+  params: { [K in keyof TParams]: ParamResolvable<TParams[K], TParams, S, TData> },
+  store: QueryStore<TData, TParams, S>
+): ResolvedParamsResult<TParams> {
   const directValues: Partial<TParams> = {};
   const paramAttachVals: Partial<Record<keyof TParams, AttachValue<unknown>>> = {};
   const resolvedParams = {} as TParams;
 
   for (const key in params) {
     const param = params[key];
-    if (isParamFn(param)) {
-      const attachVal = param($);
+    if (typeof param === 'function') {
+      const attachVal = param($, store);
       resolvedParams[key] = attachVal.value as TParams[typeof key];
       paramAttachVals[key] = attachVal;
     } else {
diff --git a/src/state/internal/signal.ts b/src/state/internal/signal.ts
index cea47c31da1..62ff4dcf641 100644
--- a/src/state/internal/signal.ts
+++ b/src/state/internal/signal.ts
@@ -25,6 +25,16 @@ export type Subscribe = (callback: () => void) => Unsubscribe;
 export type GetValue = () => unknown;
 export type SetValue = (path: unknown[], value: unknown) => void;
 
+export function $<T>(store: StoreApi<T>): AttachValue<T>;
+export function $<T, S>(store: StoreApi<T>, selector: (state: T) => S, equalityFn?: (a: S, b: S) => boolean): AttachValue<S>;
+export function $(
+  store: StoreApi<unknown>,
+  selector: (state: unknown) => unknown = identity,
+  equalityFn: (a: unknown, b: unknown) => boolean = Object.is
+) {
+  return getOrCreateAttachValue(store, selector, equalityFn);
+}
+
 const identity = <T>(x: T): T => x;
 
 const updateValue = <T>(obj: T, path: unknown[], value: unknown): T => {
@@ -153,13 +163,3 @@ function getOrCreateAttachValue<T, S>(store: StoreApi<T>, selector: (state: T) =
   byEqFn.set(equalityFn as (a: unknown, b: unknown) => boolean, rootVal);
   return rootVal as AttachValue<S>;
 }
-
-export function $<T>(store: StoreApi<T>): AttachValue<T>;
-export function $<T, S>(store: StoreApi<T>, selector: (state: T) => S, equalityFn?: (a: S, b: S) => boolean): AttachValue<S>;
-export function $(
-  store: StoreApi<unknown>,
-  selector: (state: unknown) => unknown = identity,
-  equalityFn: (a: unknown, b: unknown) => boolean = Object.is
-) {
-  return getOrCreateAttachValue(store, selector, equalityFn);
-}
diff --git a/src/state/internal/tests/QueryStoreTest.tsx b/src/state/internal/tests/QueryStoreTest.tsx
index 14f3c39a708..e08df743556 100644
--- a/src/state/internal/tests/QueryStoreTest.tsx
+++ b/src/state/internal/tests/QueryStoreTest.tsx
@@ -14,13 +14,12 @@ import { createRainbowStore } from '../createRainbowStore';
 
 const ENABLE_LOGS = false;
 
-type AddressStore = {
-  address: Address;
+type CurrencyStore = {
   currency: SupportedCurrencyKey;
-  nestedAddressTest: {
-    address: Address;
+  nestedParamTest: {
+    currency: SupportedCurrencyKey;
   };
-  setAddress: (address: Address) => void;
+  setCurrency: (currency: SupportedCurrencyKey) => void;
 };
 
 const testAddresses: Address[] = [
@@ -29,20 +28,20 @@ const testAddresses: Address[] = [
   '0x17cd072cBd45031EFc21Da538c783E0ed3b25DCc',
 ];
 
-const useAddressStore = createRainbowStore<AddressStore>((set, get) => ({
-  address: testAddresses[0],
+const useCurrencyStore = createRainbowStore<CurrencyStore>((set, get) => ({
   currency: 'USD',
-  nestedAddressTest: { address: testAddresses[0] },
+  nestedParamTest: { currency: 'USD' },
 
-  setAddress: (address: Address) => {
-    set({ address });
-    if (ENABLE_LOGS) console.log('[👤 useAddressStore 👤] New address set:', get().address);
+  setCurrency: (currency: SupportedCurrencyKey) => {
+    set({ currency });
+    if (ENABLE_LOGS) console.log('[👤 useCurrencyStore 👤] New currency set:', get().currency);
   },
 }));
 
 type TestStore = {
+  address: Address;
   userAssets: ParsedAssetsDictByChain;
-  getHighestValueAsset: () => number;
+  setAddress: (address: Address) => void;
   setUserAssets: (data: ParsedAssetsDictByChain) => void;
 };
 type QueryParams = { address: Address; currency: SupportedCurrencyKey };
@@ -69,20 +68,16 @@ export const useUserAssetsTestStore = createQueryStore<ParsedAssetsDictByChain,
     setData: (data, set) => set({ userAssets: data }),
 
     params: {
-      address: $ => $(useAddressStore).address,
-      currency: $ => $(useAddressStore).currency,
+      address: ($, store) => $(store).address,
+      currency: $ => $(useCurrencyStore).currency,
     },
     staleTime: time.minutes(1),
   },
 
-  (set, get) => ({
+  set => ({
+    address: testAddresses[0],
     userAssets: [],
-
-    getHighestValueAsset: () =>
-      Object.values(get().userAssets)
-        .flatMap(chainAssets => Object.values(chainAssets))
-        .reduce((max, asset) => Math.max(max, Number(asset.balance.display)), 0),
-
+    setAddress: (address: Address) => set({ address }),
     setUserAssets: (data: ParsedAssetsDictByChain) => set({ userAssets: data }),
   }),
 
@@ -115,16 +110,16 @@ export const UserAssetsTest = memo(function UserAssetsTest() {
         <View style={styles.buttonGroup}>
           <ButtonPressAnimation
             onPress={() => {
-              const currentAddress = useAddressStore.getState().address;
+              const currentAddress = useUserAssetsTestStore.getState().address;
               switch (currentAddress) {
                 case testAddresses[0]:
-                  useAddressStore.getState().setAddress(testAddresses[1]);
+                  useUserAssetsTestStore.getState().setAddress(testAddresses[1]);
                   break;
                 case testAddresses[1]:
-                  useAddressStore.getState().setAddress(testAddresses[2]);
+                  useUserAssetsTestStore.getState().setAddress(testAddresses[2]);
                   break;
                 case testAddresses[2]:
-                  useAddressStore.getState().setAddress(testAddresses[0]);
+                  useUserAssetsTestStore.getState().setAddress(testAddresses[0]);
                   break;
               }
             }}

From 11f8fdf8ee585883cd65c1947b8a68b78935af7d Mon Sep 17 00:00:00 2001
From: Christian Baroni <7061887+christianbaroni@users.noreply.github.com>
Date: Fri, 20 Dec 2024 04:58:02 +0000
Subject: [PATCH 17/29] [createRainbowStore] Catch up with develop

---
 src/state/internal/createRainbowStore.ts | 11 ++++++++---
 1 file changed, 8 insertions(+), 3 deletions(-)

diff --git a/src/state/internal/createRainbowStore.ts b/src/state/internal/createRainbowStore.ts
index 195d167e8e2..6f6ef4f05f5 100644
--- a/src/state/internal/createRainbowStore.ts
+++ b/src/state/internal/createRainbowStore.ts
@@ -22,6 +22,12 @@ export interface RainbowPersistConfig<S> {
    * This function will be called when persisted state versions mismatch with the one specified here.
    */
   migrate?: (persistedState: unknown, version: number) => S | Promise<S>;
+  /**
+   * A function returning another (optional) function.
+   * The main function will be called before the state rehydration.
+   * The returned function will be called after the state rehydration or when an error occurred.
+   */
+  onRehydrateStorage?: PersistOptions<S, Partial<S>>['onRehydrateStorage'];
   /**
    * A function that determines which parts of the state should be persisted.
    * By default, the entire state is persisted.
@@ -54,9 +60,7 @@ export function createRainbowStore<S = unknown>(
   createState: StateCreator<S, [], [['zustand/subscribeWithSelector', never]]>,
   persistConfig?: RainbowPersistConfig<S>
 ) {
-  if (!persistConfig) {
-    return create<S>()(subscribeWithSelector(createState));
-  }
+  if (!persistConfig) return create<S>()(subscribeWithSelector(createState));
 
   const { persistStorage, version } = createPersistStorage(persistConfig);
 
@@ -65,6 +69,7 @@ export function createRainbowStore<S = unknown>(
       persist(createState, {
         migrate: persistConfig.migrate,
         name: persistConfig.storageKey,
+        onRehydrateStorage: persistConfig.onRehydrateStorage,
         partialize: persistConfig.partialize || omitStoreMethods,
         storage: persistStorage,
         version,

From 61e97f4d56b49a6820805787a24556fb0cb2c93b Mon Sep 17 00:00:00 2001
From: Christian Baroni <7061887+christianbaroni@users.noreply.github.com>
Date: Fri, 20 Dec 2024 05:19:49 +0000
Subject: [PATCH 18/29] Remove unnecessary assignment

---
 src/state/internal/createQueryStore.ts | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/src/state/internal/createQueryStore.ts b/src/state/internal/createQueryStore.ts
index 463a1970a5b..5bf6ab94c74 100644
--- a/src/state/internal/createQueryStore.ts
+++ b/src/state/internal/createQueryStore.ts
@@ -370,10 +370,9 @@ export function createQueryStore<
 
   const createState: StateCreator<S, [], [['zustand/subscribeWithSelector', never]]> = (set, get, api) => {
     const pruneCache = (state: S): S => {
-      const now = Date.now();
       const newCache: Record<string, CacheEntry<TData>> = {};
       Object.entries(state.queryCache).forEach(([key, entry]) => {
-        if (entry && now - entry.lastFetchedAt <= cacheTime) {
+        if (entry && Date.now() - entry.lastFetchedAt <= cacheTime) {
           newCache[key] = entry;
         }
       });

From fda20e21219532ee310b9f7d758baaaa302bb15a Mon Sep 17 00:00:00 2001
From: Christian Baroni <7061887+christianbaroni@users.noreply.github.com>
Date: Fri, 20 Dec 2024 05:33:58 +0000
Subject: [PATCH 19/29] Minor docs cleanup

---
 src/state/internal/createQueryStore.ts | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/state/internal/createQueryStore.ts b/src/state/internal/createQueryStore.ts
index 5bf6ab94c74..b1289ede63f 100644
--- a/src/state/internal/createQueryStore.ts
+++ b/src/state/internal/createQueryStore.ts
@@ -119,13 +119,13 @@ export interface QueryStore<TData, TParams extends Record<string, unknown>, S ex
    */
   getStatus: () => QueryStatusInfo;
   /**
-   * Determines if the current data is expired, meaning it has exceeded the `cacheTime` duration.
+   * Determines if the current data is expired based on whether `cacheTime` has been exceeded.
    * @param override - An optional override for the default cache time, in milliseconds.
    * @returns `true` if the data is expired, otherwise `false`.
    */
   isDataExpired: (override?: number) => boolean;
   /**
-   * Determines if the current data is stale, meaning it has exceeded the `staleTime` duration.
+   * Determines if the current data is stale based on whether `staleTime` has been exceeded.
    * Stale data may be refreshed automatically in the background.
    * @param override - An optional override for the default stale time, in milliseconds.
    * @returns `true` if the data is stale, otherwise `false`.

From 96a12a74585ef02dea8fadc1b9965c3c54a34b9f Mon Sep 17 00:00:00 2001
From: Christian Baroni <7061887+christianbaroni@users.noreply.github.com>
Date: Fri, 20 Dec 2024 07:07:50 +0000
Subject: [PATCH 20/29] Split types to better obscure private internal state

---
 src/state/internal/createQueryStore.ts | 33 ++++++++++++++++----------
 1 file changed, 21 insertions(+), 12 deletions(-)

diff --git a/src/state/internal/createQueryStore.ts b/src/state/internal/createQueryStore.ts
index b1289ede63f..b73510ccc8f 100644
--- a/src/state/internal/createQueryStore.ts
+++ b/src/state/internal/createQueryStore.ts
@@ -88,8 +88,11 @@ interface CacheEntry<TData> {
  * - **`isStale(override?)`**: Checks if the current data is stale based on `staleTime`.
  * - **`reset()`**: Resets the store to its initial state, clearing data and errors.
  */
-export interface QueryStore<TData, TParams extends Record<string, unknown>, S extends StoreState<TData, TParams>>
-  extends UseBoundStore<StoreApi<S>> {
+export interface QueryStore<
+  TData,
+  TParams extends Record<string, unknown>,
+  S extends Omit<StoreState<TData, TParams>, keyof PrivateStoreState>,
+> extends UseBoundStore<StoreApi<S>> {
   /**
    * Indicates whether the store should actively fetch data.
    * When `false`, the store won't automatically refetch data.
@@ -138,8 +141,15 @@ export interface QueryStore<TData, TParams extends Record<string, unknown>, S ex
 }
 
 /**
- * The state structure managed by the query store, including query-related fields and actions.
- * This type is generally internal, and extended by user-defined states when creating a store.
+ * The private state managed by the query store, omitted from the store's public interface.
+ */
+type PrivateStoreState = {
+  subscriptionCount: number;
+};
+
+/**
+ * The full state structure managed by the query store. This type is generally internal,
+ * though the state it defines can be accessed via the store's public interface.
  */
 type StoreState<TData, TParams extends Record<string, unknown>> = Pick<
   QueryStore<TData, TParams, StoreState<TData, TParams>>,
@@ -149,7 +159,6 @@ type StoreState<TData, TParams extends Record<string, unknown>> = Pick<
   lastFetchedAt: number | null;
   queryCache: Record<string, CacheEntry<TData> | undefined>;
   status: QueryStatus;
-  subscriptionCount: number;
 };
 
 /**
@@ -263,7 +272,7 @@ interface ResolvedParamsResult<TParams> {
 /**
  * The keys that make up the internal state of the store.
  */
-type InternalStateKeys = keyof StoreState<unknown, Record<string, unknown>>;
+type InternalStateKeys = keyof (StoreState<unknown, Record<string, unknown>> & PrivateStoreState);
 
 const [persist, discard] = [true, false];
 
@@ -308,13 +317,13 @@ export function createQueryStore<
   U = unknown,
   TData = TQueryFnData,
 >(
-  config: RainbowQueryStoreConfig<TQueryFnData, TParams, TData, StoreState<TData, TParams> & U> & {
-    params?: { [K in keyof TParams]: ParamResolvable<TParams[K], TParams, StoreState<TData, TParams> & U, TData> };
+  config: RainbowQueryStoreConfig<TQueryFnData, TParams, TData, StoreState<TData, TParams> & PrivateStoreState & U> & {
+    params?: { [K in keyof TParams]: ParamResolvable<TParams[K], TParams, StoreState<TData, TParams> & PrivateStoreState & U, TData> };
   },
   customStateCreator?: StateCreator<U, [], [['zustand/subscribeWithSelector', never]]>,
-  persistConfig?: RainbowPersistConfig<StoreState<TData, TParams> & U>
+  persistConfig?: RainbowPersistConfig<StoreState<TData, TParams> & PrivateStoreState & U>
 ): QueryStore<TData, TParams, StoreState<TData, TParams> & U> {
-  type S = StoreState<TData, TParams> & U;
+  type S = StoreState<TData, TParams> & PrivateStoreState & U;
 
   const {
     fetcher,
@@ -669,8 +678,8 @@ export function createQueryStore<
     : undefined;
 
   const baseStore = persistConfig?.storageKey
-    ? createRainbowStore<StoreState<TData, TParams> & U>(createState, combinedPersistConfig)
-    : create<StoreState<TData, TParams> & U>()(subscribeWithSelector(createState));
+    ? createRainbowStore<StoreState<TData, TParams> & PrivateStoreState & U>(createState, combinedPersistConfig)
+    : create<StoreState<TData, TParams> & PrivateStoreState & U>()(subscribeWithSelector(createState));
 
   const queryCapableStore: QueryStore<TData, TParams, S> = Object.assign(baseStore, {
     enabled,

From a1d1b45c97965384012032fa1d33e634946f6111 Mon Sep 17 00:00:00 2001
From: Christian Baroni <7061887+christianbaroni@users.noreply.github.com>
Date: Sat, 21 Dec 2024 21:06:56 +0000
Subject: [PATCH 21/29] Fix cache overwriting, expose and track queryKey,
 improve `setData` docs, add overloads, misc. fixes

---
 src/state/internal/createQueryStore.ts      | 149 +++++++++----
 src/state/internal/tests/QueryStoreTest.tsx | 229 +++++++++++++-------
 2 files changed, 251 insertions(+), 127 deletions(-)

diff --git a/src/state/internal/createQueryStore.ts b/src/state/internal/createQueryStore.ts
index b73510ccc8f..6d7bf4dfa2a 100644
--- a/src/state/internal/createQueryStore.ts
+++ b/src/state/internal/createQueryStore.ts
@@ -98,6 +98,10 @@ export interface QueryStore<
    * When `false`, the store won't automatically refetch data.
    */
   enabled: boolean;
+  /**
+   * The current query key, which is a string representation of the current query parameter values.
+   */
+  queryKey: string;
   /**
    * Initiates a data fetch for the given parameters. If no parameters are provided, the store's
    * current parameters are used.
@@ -153,7 +157,7 @@ type PrivateStoreState = {
  */
 type StoreState<TData, TParams extends Record<string, unknown>> = Pick<
   QueryStore<TData, TParams, StoreState<TData, TParams>>,
-  'enabled' | 'fetch' | 'getData' | 'getStatus' | 'isDataExpired' | 'isStale' | 'reset'
+  'enabled' | 'queryKey' | 'fetch' | 'getData' | 'getStatus' | 'isDataExpired' | 'isStale' | 'reset'
 > & {
   error: Error | null;
   lastFetchedAt: number | null;
@@ -193,9 +197,23 @@ export type RainbowQueryStoreConfig<TQueryFnData, TParams extends Record<string,
   onFetched?: (data: TData, set: (partial: S | Partial<S> | ((state: S) => S | Partial<S>)) => void) => void;
   /**
    * A function that overrides the default behavior of setting the fetched data in the store's query cache.
-   * Receives the transformed data and the store's set function.
+   * Receives an object containing the transformed data, the query parameters, the query key, and the store's set function.
+   *
+   * When using `setData`, it’s important to note that you are taking full responsibility for managing query data. if your
+   * query supports variable parameters (and thus multiple query keys) and you want to cache data for each key, you’ll need
+   * to manually handle storing data based on the provided `params` or `queryKey`. Naturally, you will also bear
+   * responsibility for pruning this data in the event you do not want it persisted indefinitely.
+   *
+   * Automatic refetching per your specified `staleTime` is still managed internally by the store. While no query *data*
+   * will be cached internally if `setData` is provided, metadata such as the last fetch time for each query key is still
+   * cached and tracked by the store, unless caching is fully disabled via `disableCache: true`.
    */
-  setData?: (data: TData, set: (partial: S | Partial<S> | ((state: S) => S | Partial<S>)) => void) => void;
+  setData?: (info: {
+    data: TData;
+    params: TParams;
+    queryKey: string;
+    set: (partial: S | Partial<S> | ((state: S) => S | Partial<S>)) => void;
+  }) => void;
   /**
    * Suppresses warnings in the event a `staleTime` under the minimum is desired.
    * @default false
@@ -282,6 +300,7 @@ const SHOULD_PERSIST_INTERNAL_STATE_MAP: Record<string, boolean> = {
   error: persist,
   lastFetchedAt: persist,
   queryCache: persist,
+  queryKey: persist,
   status: persist,
 
   /* Internal state and methods to discard */
@@ -304,6 +323,31 @@ export const time = {
 
 const MIN_STALE_TIME = time.seconds(5);
 
+export function createQueryStore<
+  TQueryFnData,
+  TParams extends Record<string, unknown> = Record<string, unknown>,
+  U = unknown,
+  TData = TQueryFnData,
+>(
+  config: RainbowQueryStoreConfig<TQueryFnData, TParams, TData, StoreState<TData, TParams> & PrivateStoreState & U> & {
+    params?: { [K in keyof TParams]: ParamResolvable<TParams[K], TParams, StoreState<TData, TParams> & PrivateStoreState & U, TData> };
+  },
+  customStateCreator: StateCreator<U, [], [['zustand/subscribeWithSelector', never]]>,
+  persistConfig?: RainbowPersistConfig<StoreState<TData, TParams> & PrivateStoreState & U>
+): QueryStore<TData, TParams, StoreState<TData, TParams> & U>;
+
+export function createQueryStore<
+  TQueryFnData,
+  TParams extends Record<string, unknown> = Record<string, unknown>,
+  U = unknown,
+  TData = TQueryFnData,
+>(
+  config: RainbowQueryStoreConfig<TQueryFnData, TParams, TData, StoreState<TData, TParams> & PrivateStoreState & U> & {
+    params?: { [K in keyof TParams]: ParamResolvable<TParams[K], TParams, StoreState<TData, TParams> & PrivateStoreState & U, TData> };
+  },
+  persistConfig?: RainbowPersistConfig<StoreState<TData, TParams> & PrivateStoreState & U>
+): QueryStore<TData, TParams, StoreState<TData, TParams> & U>;
+
 /**
  * Creates a query-enabled Rainbow store with data fetching capabilities.
  * @template TQueryFnData - The raw data type returned by the fetcher
@@ -320,11 +364,17 @@ export function createQueryStore<
   config: RainbowQueryStoreConfig<TQueryFnData, TParams, TData, StoreState<TData, TParams> & PrivateStoreState & U> & {
     params?: { [K in keyof TParams]: ParamResolvable<TParams[K], TParams, StoreState<TData, TParams> & PrivateStoreState & U, TData> };
   },
-  customStateCreator?: StateCreator<U, [], [['zustand/subscribeWithSelector', never]]>,
-  persistConfig?: RainbowPersistConfig<StoreState<TData, TParams> & PrivateStoreState & U>
+  arg1?:
+    | StateCreator<U, [], [['zustand/subscribeWithSelector', never]]>
+    | RainbowPersistConfig<StoreState<TData, TParams> & PrivateStoreState & U>,
+  arg2?: RainbowPersistConfig<StoreState<TData, TParams> & PrivateStoreState & U>
 ): QueryStore<TData, TParams, StoreState<TData, TParams> & U> {
   type S = StoreState<TData, TParams> & PrivateStoreState & U;
 
+  /* If arg1 is a function, it's the customStateCreator; otherwise, it's the persistConfig. */
+  const customStateCreator = typeof arg1 === 'function' ? arg1 : () => ({}) as U;
+  const persistConfig = typeof arg1 === 'object' ? arg1 : arg2;
+
   const {
     fetcher,
     onFetched,
@@ -361,6 +411,7 @@ export function createQueryStore<
     error: null,
     lastFetchedAt: null,
     queryCache: {},
+    queryKey: '',
     status: QueryStatuses.Idle,
     subscriptionCount: 0,
   };
@@ -395,24 +446,26 @@ export function createQueryStore<
         clearTimeout(activeRefetchTimeout);
         activeRefetchTimeout = null;
       }
-      const currentQueryKey = getQueryKey(params);
+      const currentQueryKey = get().queryKey;
       const lastFetchedAt =
         get().queryCache[currentQueryKey]?.lastFetchedAt || (disableCache && lastFetchKey === currentQueryKey ? get().lastFetchedAt : null);
       const timeUntilRefetch = lastFetchedAt ? effectiveStaleTime - (Date.now() - lastFetchedAt) : effectiveStaleTime;
 
       activeRefetchTimeout = setTimeout(() => {
-        if (get().subscriptionCount > 0) {
-          get().fetch(params, { force: true });
+        const { enabled, fetch, subscriptionCount } = get();
+        if (enabled && subscriptionCount > 0) {
+          fetch(params, { force: true });
         }
       }, timeUntilRefetch);
     };
 
     const baseMethods = {
       async fetch(params: TParams | undefined, options: FetchOptions | undefined) {
-        if (!get().enabled) return;
+        if (!options?.force && !get().enabled) return;
+
         const effectiveParams = params ?? getCurrentResolvedParams();
-        const currentQueryKey = getQueryKey(effectiveParams);
-        const isLoading = get().status === QueryStatuses.Loading;
+        const { queryKey: currentQueryKey, status } = get();
+        const isLoading = status === QueryStatuses.Loading;
 
         if (activeFetchPromise && lastFetchKey === currentQueryKey && isLoading && !options?.force) {
           return activeFetchPromise;
@@ -421,6 +474,7 @@ export function createQueryStore<
         if (!options?.force && !disableCache) {
           const { errorInfo, lastFetchedAt } = get().queryCache[currentQueryKey] ?? {};
           const errorRetriesExhausted = errorInfo && errorInfo.retryCount >= maxRetries;
+
           if ((!errorInfo || errorRetriesExhausted) && lastFetchedAt && Date.now() - lastFetchedAt <= (options?.staleTime ?? staleTime)) {
             return;
           }
@@ -460,11 +514,17 @@ export function createQueryStore<
                   },
                 };
               } else if (setData) {
-                setData(transformedData, (partial: S | Partial<S> | ((state: S) => S | Partial<S>)) => {
-                  newState = typeof partial === 'function' ? { ...newState, ...partial(newState) } : { ...newState, ...partial };
+                setData({
+                  data: transformedData,
+                  params: effectiveParams,
+                  queryKey: currentQueryKey,
+                  set: (partial: S | Partial<S> | ((state: S) => S | Partial<S>)) => {
+                    newState = typeof partial === 'function' ? { ...newState, ...partial(newState) } : { ...newState, ...partial };
+                  },
                 });
                 if (!disableCache) {
                   newState.queryCache = {
+                    ...newState.queryCache,
                     [currentQueryKey]: {
                       data: null,
                       errorInfo: null,
@@ -498,12 +558,13 @@ export function createQueryStore<
             onError?.(typedError, currentRetryCount);
 
             if (currentRetryCount < maxRetries) {
-              const errorRetryDelay = typeof retryDelay === 'function' ? retryDelay(currentRetryCount, typedError) : retryDelay;
-
               if (get().subscriptionCount > 0) {
+                const errorRetryDelay = typeof retryDelay === 'function' ? retryDelay(currentRetryCount, typedError) : retryDelay;
+
                 activeRefetchTimeout = setTimeout(() => {
-                  if (get().subscriptionCount > 0) {
-                    get().fetch(effectiveParams, { force: true });
+                  const { enabled, fetch, subscriptionCount } = get();
+                  if (enabled && subscriptionCount > 0) {
+                    fetch(params, { force: true });
                   }
                 }, errorRetryDelay);
               }
@@ -555,16 +616,14 @@ export function createQueryStore<
 
       getData(params?: TParams) {
         if (disableCache) return null;
-        const currentQueryKey = getQueryKey(params ?? getCurrentResolvedParams());
+        const currentQueryKey = params ? getQueryKey(params) : get().queryKey;
         return get().queryCache[currentQueryKey]?.data ?? null;
       },
 
       getStatus() {
-        const status = get().status;
-        const currentQueryKey = getQueryKey(getCurrentResolvedParams());
+        const { queryKey, status } = get();
         const lastFetchedAt =
-          get().queryCache[currentQueryKey]?.lastFetchedAt ||
-          (disableCache && lastFetchKey === currentQueryKey ? get().lastFetchedAt : null);
+          get().queryCache[queryKey]?.lastFetchedAt || (disableCache && lastFetchKey === queryKey ? get().lastFetchedAt : null);
 
         return {
           isError: status === QueryStatuses.Error,
@@ -576,10 +635,9 @@ export function createQueryStore<
       },
 
       isDataExpired(cacheTimeOverride?: number) {
-        const currentQueryKey = getQueryKey(getCurrentResolvedParams());
+        const { queryKey } = get();
         const lastFetchedAt =
-          get().queryCache[currentQueryKey]?.lastFetchedAt ||
-          (disableCache && lastFetchKey === currentQueryKey ? get().lastFetchedAt : null);
+          get().queryCache[queryKey]?.lastFetchedAt || (disableCache && lastFetchKey === queryKey ? get().lastFetchedAt : null);
 
         if (!lastFetchedAt) return true;
         const effectiveCacheTime = cacheTimeOverride ?? cacheTime;
@@ -587,10 +645,9 @@ export function createQueryStore<
       },
 
       isStale(staleTimeOverride?: number) {
-        const currentQueryKey = getQueryKey(getCurrentResolvedParams());
+        const { queryKey } = get();
         const lastFetchedAt =
-          get().queryCache[currentQueryKey]?.lastFetchedAt ||
-          (disableCache && lastFetchKey === currentQueryKey ? get().lastFetchedAt : null);
+          get().queryCache[queryKey]?.lastFetchedAt || (disableCache && lastFetchKey === queryKey ? get().lastFetchedAt : null);
 
         if (!lastFetchedAt) return true;
         const effectiveStaleTime = staleTimeOverride ?? staleTime;
@@ -604,7 +661,7 @@ export function createQueryStore<
         }
         activeFetchPromise = null;
         lastFetchKey = null;
-        set(state => ({ ...state, ...initialData }));
+        set(state => ({ ...state, ...initialData, queryKey: getQueryKey(getCurrentResolvedParams()) }));
       },
     };
 
@@ -618,7 +675,7 @@ export function createQueryStore<
         if (state.enabled !== prevState.enabled) {
           if (state.enabled) {
             const currentParams = getCurrentResolvedParams();
-            const currentKey = getQueryKey(currentParams);
+            const currentKey = state.queryKey;
             if (currentKey !== lastFetchKey) {
               state.fetch(currentParams, { force: true });
             } else if (!state.queryCache[currentKey] || state.isStale()) {
@@ -635,12 +692,14 @@ export function createQueryStore<
         }
       });
 
-      const { fetch, isStale } = get();
+      const { enabled, fetch, isStale, queryKey } = get();
+
       const currentParams = getCurrentResolvedParams();
+      set(state => ({ ...state, queryKey: getQueryKey(currentParams) }));
 
-      if (!get().queryCache[getQueryKey(currentParams)] || isStale()) {
-        fetch(currentParams, { force: true });
-      } else {
+      if (!get().queryCache[queryKey] || isStale()) {
+        fetch(currentParams);
+      } else if (enabled) {
         scheduleNextFetch(currentParams, undefined);
       }
 
@@ -678,22 +737,23 @@ export function createQueryStore<
     : undefined;
 
   const baseStore = persistConfig?.storageKey
-    ? createRainbowStore<StoreState<TData, TParams> & PrivateStoreState & U>(createState, combinedPersistConfig)
-    : create<StoreState<TData, TParams> & PrivateStoreState & U>()(subscribeWithSelector(createState));
+    ? createRainbowStore(createState, combinedPersistConfig)
+    : create(subscribeWithSelector(createState));
 
   const queryCapableStore: QueryStore<TData, TParams, S> = Object.assign(baseStore, {
     enabled,
-    destroy: () => {
-      for (const unsub of paramUnsubscribes) unsub();
-      paramUnsubscribes = [];
-      queryCapableStore.getState().reset();
-    },
+    queryKey: baseStore.getState().queryKey,
     fetch: (params?: TParams, options?: FetchOptions) => baseStore.getState().fetch(params, options),
     getData: () => baseStore.getState().getData(),
     getStatus: () => baseStore.getState().getStatus(),
     isDataExpired: (override?: number) => baseStore.getState().isDataExpired(override),
     isStale: (override?: number) => baseStore.getState().isStale(override),
-    reset: () => baseStore.getState().reset(),
+    reset: () => {
+      for (const unsub of paramUnsubscribes) unsub();
+      paramUnsubscribes = [];
+      queryCapableStore.getState().reset();
+      queryCapableStore.setState(state => ({ ...state, ...initialData, queryKey: getQueryKey(getCurrentResolvedParams()) }));
+    },
   });
 
   if (params) {
@@ -703,8 +763,9 @@ export function createQueryStore<
   }
 
   const onParamChange = () => {
-    const newParams = getCurrentResolvedParams();
-    queryCapableStore.fetch(newParams, { force: true });
+    const newQueryKey = getQueryKey(getCurrentResolvedParams());
+    queryCapableStore.setState(state => ({ ...state, queryKey: newQueryKey }));
+    queryCapableStore.fetch();
   };
 
   let paramUnsubscribes: Unsubscribe[] = [];
diff --git a/src/state/internal/tests/QueryStoreTest.tsx b/src/state/internal/tests/QueryStoreTest.tsx
index e08df743556..7df6f19623a 100644
--- a/src/state/internal/tests/QueryStoreTest.tsx
+++ b/src/state/internal/tests/QueryStoreTest.tsx
@@ -1,95 +1,79 @@
 // ⚠️ Uncomment everything below to experiment with the QueryStore creator
 // TODO: Comment out test code below before merging
 
-import React, { memo, useEffect } from 'react';
+import React, { memo, useEffect, useMemo } from 'react';
 import { StyleSheet, View } from 'react-native';
 import { Address } from 'viem';
 import ButtonPressAnimation from '@/components/animations/ButtonPressAnimation';
-import { Text } from '@/design-system';
+import { ImgixImage } from '@/components/images';
+import { Text, useForegroundColor } from '@/design-system';
+import { logger, RainbowError } from '@/logger';
 import { SupportedCurrencyKey } from '@/references';
-import { queryUserAssets } from '@/__swaps__/screens/Swap/resources/assets/userAssets';
+import { addysHttp } from '@/resources/addys/claimables/query';
+import { parseUserAssets } from '@/__swaps__/screens/Swap/resources/assets/userAssets';
 import { ParsedAssetsDictByChain } from '@/__swaps__/types/assets';
+import { AddressAssetsReceivedMessage } from '@/__swaps__/types/refraction';
+import { useBackendNetworksStore } from '@/state/backendNetworks/backendNetworks';
 import { createQueryStore, time } from '../createQueryStore';
 import { createRainbowStore } from '../createRainbowStore';
 
 const ENABLE_LOGS = false;
 
 type CurrencyStore = {
-  currency: SupportedCurrencyKey;
   nestedParamTest: {
     currency: SupportedCurrencyKey;
   };
   setCurrency: (currency: SupportedCurrencyKey) => void;
 };
 
-const testAddresses: Address[] = [
-  '0x2e67869829c734ac13723A138a952F7A8B56e774',
-  '0xCFB83E14AEd465c79F3F82f4cfF4ff7965897644',
-  '0x17cd072cBd45031EFc21Da538c783E0ed3b25DCc',
-];
-
 const useCurrencyStore = createRainbowStore<CurrencyStore>((set, get) => ({
-  currency: 'USD',
   nestedParamTest: { currency: 'USD' },
-
   setCurrency: (currency: SupportedCurrencyKey) => {
-    set({ currency });
-    if (ENABLE_LOGS) console.log('[👤 useCurrencyStore 👤] New currency set:', get().currency);
+    set({ nestedParamTest: { currency } });
+    if (ENABLE_LOGS) console.log('[👤 useCurrencyStore 👤] New currency set:', get().nestedParamTest.currency);
   },
 }));
 
-type TestStore = {
+type UserAssetsTestStore = {
   address: Address;
-  userAssets: ParsedAssetsDictByChain;
   setAddress: (address: Address) => void;
-  setUserAssets: (data: ParsedAssetsDictByChain) => void;
 };
-type QueryParams = { address: Address; currency: SupportedCurrencyKey };
 
-function logFetchInfo(params: QueryParams) {
-  const formattedTimeWithSeconds = new Date(Date.now()).toLocaleTimeString('en-US', {
-    hour: '2-digit',
-    minute: '2-digit',
-    second: '2-digit',
-  });
-  console.log('[🔄 UserAssetsTest - logFetchInfo 🔄]', '\nTime:', formattedTimeWithSeconds, '\nParams:', {
-    address: params.address,
-    currency: params.currency,
-    raw: JSON.stringify(Object.values(params), null, 2),
-  });
-}
+type UserAssetsQueryParams = { address: Address; currency: SupportedCurrencyKey };
 
-export const useUserAssetsTestStore = createQueryStore<ParsedAssetsDictByChain, QueryParams, TestStore>(
-  {
-    fetcher: ({ address, currency }) => {
-      if (ENABLE_LOGS) logFetchInfo({ address, currency });
-      return queryUserAssets({ address, currency });
-    },
-    setData: (data, set) => set({ userAssets: data }),
+const testAddresses: Address[] = [
+  '0x2e67869829c734ac13723A138a952F7A8B56e774',
+  '0xCFB83E14AEd465c79F3F82f4cfF4ff7965897644',
+  '0x17cd072cBd45031EFc21Da538c783E0ed3b25DCc',
+];
 
+export const useUserAssetsTestStore = createQueryStore<ParsedAssetsDictByChain, UserAssetsQueryParams, UserAssetsTestStore>(
+  {
+    fetcher: ({ address, currency }) => simpleUserAssetsQuery({ address, currency }),
     params: {
       address: ($, store) => $(store).address,
-      currency: $ => $(useCurrencyStore).currency,
+      currency: $ => $(useCurrencyStore).nestedParamTest.currency,
     },
     staleTime: time.minutes(1),
   },
 
   set => ({
     address: testAddresses[0],
-    userAssets: [],
     setAddress: (address: Address) => set({ address }),
-    setUserAssets: (data: ParsedAssetsDictByChain) => set({ userAssets: data }),
   }),
 
-  { storageKey: 'userAssetsQueryStoreTest' }
+  { storageKey: 'queryStoreTest' }
 );
 
 export const UserAssetsTest = memo(function UserAssetsTest() {
-  const data = useUserAssetsTestStore(state => state.userAssets);
+  const data = useUserAssetsTestStore(state => state.getData());
   const enabled = useUserAssetsTestStore(state => state.enabled);
 
+  const firstFiveCoinIconUrls = useMemo(() => (data ? getFirstFiveCoinIconUrls(data) : Array.from({ length: 5 }).map(() => '')), [data]);
+  const skeletonColor = useForegroundColor('fillQuaternary');
+
   useEffect(() => {
-    if (ENABLE_LOGS) {
+    if (ENABLE_LOGS && data) {
       const first5Tokens = Object.values(data)
         .flatMap(chainAssets => Object.values(chainAssets))
         .slice(0, 5);
@@ -102,50 +86,120 @@ export const UserAssetsTest = memo(function UserAssetsTest() {
   }, [enabled]);
 
   return (
-    data && (
-      <View style={styles.container}>
-        <Text color="label" size="17pt" weight="heavy">
-          Number of assets: {Object.values(data).reduce((acc, chainAssets) => acc + Object.keys(chainAssets).length, 0)}
-        </Text>
-        <View style={styles.buttonGroup}>
-          <ButtonPressAnimation
-            onPress={() => {
-              const currentAddress = useUserAssetsTestStore.getState().address;
-              switch (currentAddress) {
-                case testAddresses[0]:
-                  useUserAssetsTestStore.getState().setAddress(testAddresses[1]);
-                  break;
-                case testAddresses[1]:
-                  useUserAssetsTestStore.getState().setAddress(testAddresses[2]);
-                  break;
-                case testAddresses[2]:
-                  useUserAssetsTestStore.getState().setAddress(testAddresses[0]);
-                  break;
-              }
-            }}
-            style={styles.button}
-          >
-            <Text color="label" size="17pt" weight="heavy">
-              Shuffle Address
-            </Text>
-          </ButtonPressAnimation>
-          <ButtonPressAnimation
-            onPress={() => {
-              useUserAssetsTestStore.setState({ enabled: !enabled });
-            }}
-            style={styles.button}
-          >
-            <Text color="label" size="17pt" weight="heavy">
-              {useUserAssetsTestStore.getState().enabled ? 'Disable Fetching' : 'Enable Fetching'}
-            </Text>
-          </ButtonPressAnimation>
-        </View>
+    <View style={styles.container}>
+      <View style={styles.coinIconContainer}>
+        {firstFiveCoinIconUrls.map((url, index) =>
+          url ? (
+            <ImgixImage enableFasterImage key={index} source={{ uri: url }} style={styles.coinIcon} />
+          ) : (
+            <View key={index} style={[styles.coinIcon, { backgroundColor: skeletonColor }]} />
+          )
+        )}
       </View>
-    )
+      <Text color="label" size="17pt" weight="heavy">
+        {data
+          ? `Number of assets: ${Object.values(data).reduce((acc, chainAssets) => acc + Object.keys(chainAssets).length, 0)}`
+          : 'Loading…'}
+      </Text>
+      <View style={styles.buttonGroup}>
+        <ButtonPressAnimation
+          onPress={() => {
+            const currentAddress = useUserAssetsTestStore.getState().address;
+            switch (currentAddress) {
+              case testAddresses[0]:
+                useUserAssetsTestStore.getState().setAddress(testAddresses[1]);
+                break;
+              case testAddresses[1]:
+                useUserAssetsTestStore.getState().setAddress(testAddresses[2]);
+                break;
+              case testAddresses[2]:
+                useUserAssetsTestStore.getState().setAddress(testAddresses[0]);
+                break;
+            }
+          }}
+          style={styles.button}
+        >
+          <Text color="label" size="17pt" weight="heavy">
+            Shuffle Address
+          </Text>
+        </ButtonPressAnimation>
+        <ButtonPressAnimation
+          onPress={() => {
+            useUserAssetsTestStore.setState({ enabled: !enabled });
+          }}
+          style={styles.button}
+        >
+          <Text color="label" size="17pt" weight="heavy">
+            {useUserAssetsTestStore.getState().enabled ? 'Disable Fetching' : 'Enable Fetching'}
+          </Text>
+        </ButtonPressAnimation>
+      </View>
+    </View>
   );
 });
 
-if (ENABLE_LOGS) console.log('[💾 UserAssetsTest 💾] Initial data exists:', !!useUserAssetsTestStore.getState().userAssets);
+if (ENABLE_LOGS) console.log('[💾 UserAssetsTest 💾] Initial data exists:', !!useUserAssetsTestStore.getState().getData());
+
+// eslint-disable-next-line @typescript-eslint/no-unused-vars
+function logFetchInfo(params: UserAssetsQueryParams) {
+  const formattedTimeWithSeconds = new Date(Date.now()).toLocaleTimeString('en-US', {
+    hour: '2-digit',
+    minute: '2-digit',
+    second: '2-digit',
+  });
+  console.log('[🔄 UserAssetsTest - logFetchInfo 🔄]', '\nTime:', formattedTimeWithSeconds, '\nParams:', {
+    address: params.address,
+    currency: params.currency,
+    raw: JSON.stringify(Object.values(params), null, 2),
+  });
+}
+
+function getFirstFiveCoinIconUrls(data: ParsedAssetsDictByChain) {
+  const result: string[] = [];
+  outer: for (const chainAssets of Object.values(data)) {
+    for (const token of Object.values(chainAssets)) {
+      if (token.icon_url) {
+        result.push(token.icon_url);
+        if (result.length === 5) {
+          break outer;
+        }
+      }
+    }
+  }
+  return result;
+}
+
+type FetchUserAssetsArgs = {
+  address: Address | string;
+  currency: SupportedCurrencyKey;
+  testnetMode?: boolean;
+};
+
+export async function simpleUserAssetsQuery({ address, currency }: FetchUserAssetsArgs): Promise<ParsedAssetsDictByChain> {
+  if (!address) return {};
+  try {
+    const url = `/${useBackendNetworksStore.getState().getSupportedChainIds().join(',')}/${address}/assets?currency=${currency.toLowerCase()}`;
+    const res = await addysHttp.get<AddressAssetsReceivedMessage>(url, {
+      timeout: time.seconds(20),
+    });
+    const chainIdsInResponse = res?.data?.meta?.chain_ids || [];
+    const assets = res?.data?.payload?.assets?.filter(asset => !asset.asset.defi_position) || [];
+
+    if (assets.length && chainIdsInResponse.length) {
+      return parseUserAssets({
+        assets,
+        chainIds: chainIdsInResponse,
+        currency,
+      });
+    }
+    return {};
+  } catch (e) {
+    logger.error(new RainbowError('[userAssetsQueryFunction]: Failed to fetch user assets'), {
+      message: (e as Error)?.message,
+    });
+    return {};
+  }
+}
 
 const styles = StyleSheet.create({
   button: {
@@ -162,6 +216,15 @@ const styles = StyleSheet.create({
     gap: 24,
     justifyContent: 'center',
   },
+  coinIcon: {
+    borderRadius: 16,
+    height: 32,
+    width: 32,
+  },
+  coinIconContainer: {
+    flexDirection: 'row',
+    gap: 12,
+  },
   container: {
     alignItems: 'center',
     flex: 1,

From 8b11f97958d406eb3b37a9c83248826fc6b28183 Mon Sep 17 00:00:00 2001
From: Christian Baroni <7061887+christianbaroni@users.noreply.github.com>
Date: Sat, 21 Dec 2024 21:25:52 +0000
Subject: [PATCH 22/29] Fix signal logging

---
 src/state/internal/signal.ts | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/state/internal/signal.ts b/src/state/internal/signal.ts
index 62ff4dcf641..2bfe57c959e 100644
--- a/src/state/internal/signal.ts
+++ b/src/state/internal/signal.ts
@@ -118,8 +118,6 @@ function getOrCreateAttachValue<T, S>(store: StoreApi<T>, selector: (state: T) =
 
   const [subscribe, getVal, setVal] = createSignal(store, selector, equalityFn);
 
-  if (ENABLE_LOGS) console.log('[🌀 AttachValue 🌀] Created root attachValue:', { selector: selector.toString() });
-
   const localCache = new Map<string, AttachValue<unknown>>();
 
   const createAttachValue = (fullPath: string): AttachValue<unknown> => {
@@ -138,6 +136,8 @@ function getOrCreateAttachValue<T, S>(store: StoreApi<T>, selector: (state: T) =
         if (cached) {
           if (ENABLE_LOGS) console.log('[🌀 AttachValue 🌀] Cache hit for:', pathKey);
           return cached;
+        } else if (ENABLE_LOGS) {
+          console.log('[🌀 AttachValue 🌀] Created root attachValue:', pathKey);
         }
         const val = createAttachValue(pathKey);
         attachValueSubscriptionMap.set(val, subscribe);

From c677ad2488f018cc1e9e305ad1dbbe920b51d038 Mon Sep 17 00:00:00 2001
From: Christian Baroni <7061887+christianbaroni@users.noreply.github.com>
Date: Sun, 22 Dec 2024 21:10:13 +0000
Subject: [PATCH 23/29] Organize types, skip string conversion if unneeded in
 getOrCreateAttachValue

---
 src/state/internal/createQueryStore.ts      | 34 ++++++++++-----------
 src/state/internal/signal.ts                |  3 +-
 src/state/internal/tests/QueryStoreTest.tsx |  8 ++---
 3 files changed, 23 insertions(+), 22 deletions(-)

diff --git a/src/state/internal/createQueryStore.ts b/src/state/internal/createQueryStore.ts
index 6d7bf4dfa2a..41ad4ad0980 100644
--- a/src/state/internal/createQueryStore.ts
+++ b/src/state/internal/createQueryStore.ts
@@ -174,17 +174,6 @@ export type RainbowQueryStoreConfig<TQueryFnData, TParams extends Record<string,
    * Receives parameters of type `TParams` and returns either a promise or a raw data value of type `TQueryFnData`.
    */
   fetcher: (params: TParams) => TQueryFnData | Promise<TQueryFnData>;
-  /**
-   * The maximum number of times to retry a failed fetch operation.
-   * @default 3
-   */
-  maxRetries?: number;
-  /**
-   * The delay between retries after a fetch error occurs, in milliseconds, defined as a number or a function that
-   * receives the error and current retry count and returns a number.
-   * @default time.seconds(5)
-   */
-  retryDelay?: number | ((retryCount: number, error: Error) => number);
   /**
    * A callback invoked whenever a fetch operation fails.
    * Receives the error and the current retry count.
@@ -199,7 +188,7 @@ export type RainbowQueryStoreConfig<TQueryFnData, TParams extends Record<string,
    * A function that overrides the default behavior of setting the fetched data in the store's query cache.
    * Receives an object containing the transformed data, the query parameters, the query key, and the store's set function.
    *
-   * When using `setData`, it’s important to note that you are taking full responsibility for managing query data. if your
+   * When using `setData`, it’s important to note that you are taking full responsibility for managing query data. If your
    * query supports variable parameters (and thus multiple query keys) and you want to cache data for each key, you’ll need
    * to manually handle storing data based on the provided `params` or `queryKey`. Naturally, you will also bear
    * responsibility for pruning this data in the event you do not want it persisted indefinitely.
@@ -214,11 +203,6 @@ export type RainbowQueryStoreConfig<TQueryFnData, TParams extends Record<string,
     queryKey: string;
     set: (partial: S | Partial<S> | ((state: S) => S | Partial<S>)) => void;
   }) => void;
-  /**
-   * Suppresses warnings in the event a `staleTime` under the minimum is desired.
-   * @default false
-   */
-  suppressStaleTimeWarning?: boolean;
   /**
    * A function to transform the raw fetched data (`TQueryFnData`) into another form (`TData`).
    * If not provided, the raw data returned by `fetcher` is used.
@@ -245,6 +229,11 @@ export type RainbowQueryStoreConfig<TQueryFnData, TParams extends Record<string,
    * @default true
    */
   enabled?: boolean;
+  /**
+   * The maximum number of times to retry a failed fetch operation.
+   * @default 3
+   */
+  maxRetries?: number;
   /**
    * Parameters to be passed to the fetcher, defined as either direct values or `ParamResolvable` functions.
    * Dynamic parameters using `AttachValue` will cause the store to refetch when their values change.
@@ -252,6 +241,12 @@ export type RainbowQueryStoreConfig<TQueryFnData, TParams extends Record<string,
   params?: {
     [K in keyof TParams]: ParamResolvable<TParams[K], TParams, S, TData>;
   };
+  /**
+   * The delay between retries after a fetch error occurs, in milliseconds, defined as a number or a function that
+   * receives the error and current retry count and returns a number.
+   * @default time.seconds(5)
+   */
+  retryDelay?: number | ((retryCount: number, error: Error) => number);
   /**
    * The duration, in milliseconds, that data is considered fresh after fetching.
    * After becoming stale, the store may automatically refetch data in the background if there are active subscribers.
@@ -260,6 +255,11 @@ export type RainbowQueryStoreConfig<TQueryFnData, TParams extends Record<string,
    * @default time.minutes(2)
    */
   staleTime?: number;
+  /**
+   * Suppresses warnings in the event a `staleTime` under the minimum is desired.
+   * @default false
+   */
+  suppressStaleTimeWarning?: boolean;
 };
 
 /**
diff --git a/src/state/internal/signal.ts b/src/state/internal/signal.ts
index 2bfe57c959e..536f0ba5690 100644
--- a/src/state/internal/signal.ts
+++ b/src/state/internal/signal.ts
@@ -131,7 +131,8 @@ function getOrCreateAttachValue<T, S>(store: StoreApi<T>, selector: (state: T) =
           }
           return v;
         }
-        const pathKey = fullPath ? `${fullPath}.${key.toString()}` : key.toString();
+        const keyString = typeof key === 'string' ? key : key.toString();
+        const pathKey = fullPath ? `${fullPath}.${keyString}` : keyString;
         const cached = localCache.get(pathKey);
         if (cached) {
           if (ENABLE_LOGS) console.log('[🌀 AttachValue 🌀] Cache hit for:', pathKey);
diff --git a/src/state/internal/tests/QueryStoreTest.tsx b/src/state/internal/tests/QueryStoreTest.tsx
index 7df6f19623a..6041d131a0b 100644
--- a/src/state/internal/tests/QueryStoreTest.tsx
+++ b/src/state/internal/tests/QueryStoreTest.tsx
@@ -96,7 +96,7 @@ export const UserAssetsTest = memo(function UserAssetsTest() {
           )
         )}
       </View>
-      <Text color="label" size="17pt" weight="heavy">
+      <Text align="center" color="label" size="17pt" weight="heavy">
         {data
           ? `Number of assets: ${Object.values(data).reduce((acc, chainAssets) => acc + Object.keys(chainAssets).length, 0)}`
           : 'Loading…'}
@@ -119,7 +119,7 @@ export const UserAssetsTest = memo(function UserAssetsTest() {
           }}
           style={styles.button}
         >
-          <Text color="label" size="17pt" weight="heavy">
+          <Text align="center" color="label" size="17pt" weight="heavy">
             Shuffle Address
           </Text>
         </ButtonPressAnimation>
@@ -129,7 +129,7 @@ export const UserAssetsTest = memo(function UserAssetsTest() {
           }}
           style={styles.button}
         >
-          <Text color="label" size="17pt" weight="heavy">
+          <Text align="center" color="label" size="17pt" weight="heavy">
             {useUserAssetsTestStore.getState().enabled ? 'Disable Fetching' : 'Enable Fetching'}
           </Text>
         </ButtonPressAnimation>
@@ -194,7 +194,7 @@ export async function simpleUserAssetsQuery({ address, currency }: FetchUserAsse
     }
     return {};
   } catch (e) {
-    logger.error(new RainbowError('[userAssetsQueryFunction]: Failed to fetch user assets'), {
+    logger.error(new RainbowError('[simpleUserAssetsQuery]: Failed to fetch user assets'), {
       message: (e as Error)?.message,
     });
     return {};

From c898f97c4d4d0642d490deddd0fc60de1bdb620e Mon Sep 17 00:00:00 2001
From: Christian Baroni <7061887+christianbaroni@users.noreply.github.com>
Date: Sun, 29 Dec 2024 22:06:25 +0000
Subject: [PATCH 24/29] Add AbortController, abortInterruptedFetches setting
 (default true)

---
 src/state/internal/createQueryStore.ts | 56 ++++++++++++++++++++++----
 1 file changed, 49 insertions(+), 7 deletions(-)

diff --git a/src/state/internal/createQueryStore.ts b/src/state/internal/createQueryStore.ts
index 41ad4ad0980..491e8a81f58 100644
--- a/src/state/internal/createQueryStore.ts
+++ b/src/state/internal/createQueryStore.ts
@@ -24,7 +24,7 @@ export const QueryStatuses = {
  * - **`'error'`**  : The most recent request encountered an error.
  * - **`'idle'`**   : No request in progress, no error, no data yet.
  * - **`'loading'`** : A request is currently in progress.
- * - **`'success'`** : The most recent request has succeeded, and `data` is available.
+ * - **`'success'`** : The most recent request has succeeded and data is available.
  */
 export type QueryStatus = (typeof QueryStatuses)[keyof typeof QueryStatuses];
 
@@ -81,6 +81,7 @@ interface CacheEntry<TData> {
  *
  * In addition to Zustand's store API (such as `getState()` and `subscribe()`), this interface provides:
  * - **`enabled`**: A boolean indicating if the store is actively fetching data.
+ * - **`queryKey`**: A string representation of the current query parameter values.
  * - **`fetch(params, options)`**: Initiates a data fetch operation.
  * - **`getData(params)`**: Returns the cached data, if available, for the current query parameters.
  * - **`getStatus()`**: Returns expanded status information for the current query parameters.
@@ -168,7 +169,7 @@ type StoreState<TData, TParams extends Record<string, unknown>> = Pick<
 /**
  * Configuration options for creating a query-enabled Rainbow store.
  */
-export type RainbowQueryStoreConfig<TQueryFnData, TParams extends Record<string, unknown>, TData, S extends StoreState<TData, TParams>> = {
+export type QueryStoreConfig<TQueryFnData, TParams extends Record<string, unknown>, TData, S extends StoreState<TData, TParams>> = {
   /**
    * A function responsible for fetching data from a remote source.
    * Receives parameters of type `TParams` and returns either a promise or a raw data value of type `TQueryFnData`.
@@ -208,6 +209,13 @@ export type RainbowQueryStoreConfig<TQueryFnData, TParams extends Record<string,
    * If not provided, the raw data returned by `fetcher` is used.
    */
   transform?: (data: TQueryFnData) => TData;
+  /**
+   * If `true`, the store will abort any partially completed fetches when:
+   * - A new fetch is initiated due to a change in parameters
+   * - All components subscribed to the store via selectors are unmounted
+   * @default true
+   */
+  abortInterruptedFetches?: boolean;
   /**
    * The maximum duration, in milliseconds, that fetched data is considered fresh.
    * After this time, data is considered expired and will be refetched when requested.
@@ -313,6 +321,8 @@ const SHOULD_PERSIST_INTERNAL_STATE_MAP: Record<string, boolean> = {
   subscriptionCount: discard,
 } satisfies Record<InternalStateKeys, boolean>;
 
+const ABORT_ERROR = new Error('[createQueryStore: AbortError] Fetch interrupted');
+
 export const time = {
   seconds: (n: number) => n * 1000,
   minutes: (n: number) => time.seconds(n * 60),
@@ -329,7 +339,7 @@ export function createQueryStore<
   U = unknown,
   TData = TQueryFnData,
 >(
-  config: RainbowQueryStoreConfig<TQueryFnData, TParams, TData, StoreState<TData, TParams> & PrivateStoreState & U> & {
+  config: QueryStoreConfig<TQueryFnData, TParams, TData, StoreState<TData, TParams> & PrivateStoreState & U> & {
     params?: { [K in keyof TParams]: ParamResolvable<TParams[K], TParams, StoreState<TData, TParams> & PrivateStoreState & U, TData> };
   },
   customStateCreator: StateCreator<U, [], [['zustand/subscribeWithSelector', never]]>,
@@ -342,7 +352,7 @@ export function createQueryStore<
   U = unknown,
   TData = TQueryFnData,
 >(
-  config: RainbowQueryStoreConfig<TQueryFnData, TParams, TData, StoreState<TData, TParams> & PrivateStoreState & U> & {
+  config: QueryStoreConfig<TQueryFnData, TParams, TData, StoreState<TData, TParams> & PrivateStoreState & U> & {
     params?: { [K in keyof TParams]: ParamResolvable<TParams[K], TParams, StoreState<TData, TParams> & PrivateStoreState & U, TData> };
   },
   persistConfig?: RainbowPersistConfig<StoreState<TData, TParams> & PrivateStoreState & U>
@@ -361,7 +371,7 @@ export function createQueryStore<
   U = unknown,
   TData = TQueryFnData,
 >(
-  config: RainbowQueryStoreConfig<TQueryFnData, TParams, TData, StoreState<TData, TParams> & PrivateStoreState & U> & {
+  config: QueryStoreConfig<TQueryFnData, TParams, TData, StoreState<TData, TParams> & PrivateStoreState & U> & {
     params?: { [K in keyof TParams]: ParamResolvable<TParams[K], TParams, StoreState<TData, TParams> & PrivateStoreState & U, TData> };
   },
   arg1?:
@@ -376,6 +386,7 @@ export function createQueryStore<
   const persistConfig = typeof arg1 === 'object' ? arg1 : arg2;
 
   const {
+    abortInterruptedFetches = true,
     fetcher,
     onFetched,
     transform,
@@ -393,7 +404,7 @@ export function createQueryStore<
 
   if (IS_DEV && !suppressStaleTimeWarning && staleTime < MIN_STALE_TIME) {
     console.warn(
-      `[RainbowQueryStore${persistConfig?.storageKey ? `: ${persistConfig.storageKey}` : ''}] ❌ Stale times under ${
+      `[createQueryStore${persistConfig?.storageKey ? `: ${persistConfig.storageKey}` : ''}] ❌ Stale times under ${
         MIN_STALE_TIME / 1000
       } seconds are not recommended.`
     );
@@ -402,6 +413,7 @@ export function createQueryStore<
   let directValues: Partial<TParams> = {};
   let paramAttachVals: Partial<Record<keyof TParams, AttachValue<unknown>>> = {};
 
+  let activeAbortController: AbortController | null = null;
   let activeFetchPromise: Promise<void> | null = null;
   let activeRefetchTimeout: NodeJS.Timeout | null = null;
   let lastFetchKey: string | null = null;
@@ -428,6 +440,30 @@ export function createQueryStore<
     return currentParams as TParams;
   };
 
+  const fetchWithAbortControl = async (params: TParams): Promise<TQueryFnData> => {
+    const abortController = new AbortController();
+    activeAbortController = abortController;
+
+    try {
+      return await new Promise((resolve, reject) => {
+        abortController.signal.addEventListener('abort', () => reject(ABORT_ERROR), { once: true });
+
+        Promise.resolve(fetcher(params)).then(resolve, reject);
+      });
+    } finally {
+      if (activeAbortController === abortController) {
+        activeAbortController = null;
+      }
+    }
+  };
+
+  const abortActiveFetch = () => {
+    if (activeAbortController) {
+      activeAbortController.abort();
+      activeAbortController = null;
+    }
+  };
+
   const createState: StateCreator<S, [], [['zustand/subscribeWithSelector', never]]> = (set, get, api) => {
     const pruneCache = (state: S): S => {
       const newCache: Record<string, CacheEntry<TData>> = {};
@@ -471,6 +507,8 @@ export function createQueryStore<
           return activeFetchPromise;
         }
 
+        if (abortInterruptedFetches) abortActiveFetch();
+
         if (!options?.force && !disableCache) {
           const { errorInfo, lastFetchedAt } = get().queryCache[currentQueryKey] ?? {};
           const errorRetriesExhausted = errorInfo && errorInfo.retryCount >= maxRetries;
@@ -485,7 +523,7 @@ export function createQueryStore<
 
         const fetchOperation = async () => {
           try {
-            const rawResult = await fetcher(effectiveParams);
+            const rawResult = await (abortInterruptedFetches ? fetchWithAbortControl(effectiveParams) : fetcher(effectiveParams));
             let transformedData: TData;
             try {
               transformedData = transform ? transform(rawResult) : (rawResult as TData);
@@ -551,6 +589,8 @@ export function createQueryStore<
               }
             }
           } catch (error) {
+            if (error === ABORT_ERROR) return;
+
             const typedError = error instanceof Error ? error : new Error(String(error));
             const entry = get().queryCache[currentQueryKey];
             const currentRetryCount = entry?.errorInfo?.retryCount ?? 0;
@@ -659,6 +699,7 @@ export function createQueryStore<
           clearTimeout(activeRefetchTimeout);
           activeRefetchTimeout = null;
         }
+        if (abortInterruptedFetches) abortActiveFetch();
         activeFetchPromise = null;
         lastFetchKey = null;
         set(state => ({ ...state, ...initialData, queryKey: getQueryKey(getCurrentResolvedParams()) }));
@@ -713,6 +754,7 @@ export function createQueryStore<
               clearTimeout(activeRefetchTimeout);
               activeRefetchTimeout = null;
             }
+            if (abortInterruptedFetches) abortActiveFetch();
           }
           return { ...state, subscriptionCount: newCount };
         });

From a6f899804b8f2c78b96324524593cc2ef23b37be Mon Sep 17 00:00:00 2001
From: Christian Baroni <7061887+christianbaroni@users.noreply.github.com>
Date: Mon, 30 Dec 2024 21:40:51 +0000
Subject: [PATCH 25/29] Ensure effectiveParams and currentQueryKey are always
 in sync within each fetch operation

---
 src/state/internal/createQueryStore.ts | 15 ++++++++-------
 1 file changed, 8 insertions(+), 7 deletions(-)

diff --git a/src/state/internal/createQueryStore.ts b/src/state/internal/createQueryStore.ts
index 491e8a81f58..0418c1a5029 100644
--- a/src/state/internal/createQueryStore.ts
+++ b/src/state/internal/createQueryStore.ts
@@ -482,7 +482,7 @@ export function createQueryStore<
         clearTimeout(activeRefetchTimeout);
         activeRefetchTimeout = null;
       }
-      const currentQueryKey = get().queryKey;
+      const currentQueryKey = getQueryKey(params);
       const lastFetchedAt =
         get().queryCache[currentQueryKey]?.lastFetchedAt || (disableCache && lastFetchKey === currentQueryKey ? get().lastFetchedAt : null);
       const timeUntilRefetch = lastFetchedAt ? effectiveStaleTime - (Date.now() - lastFetchedAt) : effectiveStaleTime;
@@ -500,8 +500,8 @@ export function createQueryStore<
         if (!options?.force && !get().enabled) return;
 
         const effectiveParams = params ?? getCurrentResolvedParams();
-        const { queryKey: currentQueryKey, status } = get();
-        const isLoading = status === QueryStatuses.Loading;
+        const currentQueryKey = getQueryKey(effectiveParams);
+        const isLoading = get().status === QueryStatuses.Loading;
 
         if (activeFetchPromise && lastFetchKey === currentQueryKey && isLoading && !options?.force) {
           return activeFetchPromise;
@@ -733,12 +733,13 @@ export function createQueryStore<
         }
       });
 
-      const { enabled, fetch, isStale, queryKey } = get();
-
+      const { enabled, fetch, isStale } = get();
       const currentParams = getCurrentResolvedParams();
-      set(state => ({ ...state, queryKey: getQueryKey(currentParams) }));
+      const currentQueryKey = getQueryKey(currentParams);
+
+      set(state => ({ ...state, queryKey: currentQueryKey }));
 
-      if (!get().queryCache[queryKey] || isStale()) {
+      if (!get().queryCache[currentQueryKey] || isStale()) {
         fetch(currentParams);
       } else if (enabled) {
         scheduleNextFetch(currentParams, undefined);

From 055cdd47699ee91e8480eac3f3e3894728590303 Mon Sep 17 00:00:00 2001
From: Christian Baroni <7061887+christianbaroni@users.noreply.github.com>
Date: Wed, 1 Jan 2025 08:19:21 +0000
Subject: [PATCH 26/29] Improve logging, fix error handling, improve param
 equality check accuracy, make fetch return data, expose params to transform,
 misc. optimizations

---
 src/state/internal/createQueryStore.ts | 107 ++++++++++++++++---------
 1 file changed, 70 insertions(+), 37 deletions(-)

diff --git a/src/state/internal/createQueryStore.ts b/src/state/internal/createQueryStore.ts
index 0418c1a5029..b2cbd119068 100644
--- a/src/state/internal/createQueryStore.ts
+++ b/src/state/internal/createQueryStore.ts
@@ -1,3 +1,4 @@
+import equal from 'react-fast-compare';
 import { StateCreator, StoreApi, UseBoundStore, create } from 'zustand';
 import { subscribeWithSelector } from 'zustand/middleware';
 import { IS_DEV } from '@/env';
@@ -5,8 +6,6 @@ import { RainbowError, logger } from '@/logger';
 import { RainbowPersistConfig, createRainbowStore, omitStoreMethods } from './createRainbowStore';
 import { $, AttachValue, SignalFunction, Unsubscribe, attachValueSubscriptionMap } from './signal';
 
-const ENABLE_LOGS = false;
-
 /**
  * A set of constants representing the various stages of a query's remote data fetching process.
  */
@@ -110,7 +109,7 @@ export interface QueryStore<
    * @param options - Optional {@link FetchOptions} to customize the fetch behavior.
    * @returns A promise that resolves when the fetch operation completes.
    */
-  fetch: (params?: TParams, options?: FetchOptions) => Promise<void>;
+  fetch: (params?: TParams, options?: FetchOptions) => Promise<TData | null>;
   /**
    * Returns the cached data, if available, for the current query params.
    * @returns The cached data, or `null` if no data is available.
@@ -208,7 +207,7 @@ export type QueryStoreConfig<TQueryFnData, TParams extends Record<string, unknow
    * A function to transform the raw fetched data (`TQueryFnData`) into another form (`TData`).
    * If not provided, the raw data returned by `fetcher` is used.
    */
-  transform?: (data: TQueryFnData) => TData;
+  transform?: (data: TQueryFnData, params: TParams) => TData;
   /**
    * If `true`, the store will abort any partially completed fetches when:
    * - A new fetch is initiated due to a change in parameters
@@ -223,11 +222,13 @@ export type QueryStoreConfig<TQueryFnData, TParams extends Record<string, unknow
    */
   cacheTime?: number;
   /**
-   * If `true`, the store's caching mechanisms will be fully disabled, meaning that the store will
-   * always refetch data on every call to `fetch()`, and the fetched data will not be stored unless
-   * a `setData` function is provided.
-   *
-   * Disable caching if you always want fresh data on refetch.
+   * If `true`, the store will log debug messages to the console.
+   * @default false
+   */
+  debugMode?: boolean;
+  /**
+   * Controls whether the store's caching mechanisms are disabled. When disabled, the store will always refetch
+   * data when params change, and fetched data will not be stored unless a `setData` function is provided.
    * @default false
    */
   disableCache?: boolean;
@@ -391,6 +392,7 @@ export function createQueryStore<
     onFetched,
     transform,
     cacheTime = time.days(7),
+    debugMode = false,
     disableCache = false,
     enabled = true,
     maxRetries = 3,
@@ -414,10 +416,12 @@ export function createQueryStore<
   let paramAttachVals: Partial<Record<keyof TParams, AttachValue<unknown>>> = {};
 
   let activeAbortController: AbortController | null = null;
-  let activeFetchPromise: Promise<void> | null = null;
+  let activeFetchPromise: Promise<TData | null> | null = null;
   let activeRefetchTimeout: NodeJS.Timeout | null = null;
   let lastFetchKey: string | null = null;
 
+  const enableLogs = IS_DEV && debugMode;
+
   const initialData = {
     enabled,
     error: null,
@@ -482,9 +486,11 @@ export function createQueryStore<
         clearTimeout(activeRefetchTimeout);
         activeRefetchTimeout = null;
       }
+
       const currentQueryKey = getQueryKey(params);
       const lastFetchedAt =
-        get().queryCache[currentQueryKey]?.lastFetchedAt || (disableCache && lastFetchKey === currentQueryKey ? get().lastFetchedAt : null);
+        (!disableCache && get().queryCache[currentQueryKey]?.lastFetchedAt) ||
+        (disableCache && lastFetchKey === currentQueryKey ? get().lastFetchedAt : null);
       const timeUntilRefetch = lastFetchedAt ? effectiveStaleTime - (Date.now() - lastFetchedAt) : effectiveStaleTime;
 
       activeRefetchTimeout = setTimeout(() => {
@@ -497,7 +503,7 @@ export function createQueryStore<
 
     const baseMethods = {
       async fetch(params: TParams | undefined, options: FetchOptions | undefined) {
-        if (!options?.force && !get().enabled) return;
+        if (!options?.force && !get().enabled) return null;
 
         const effectiveParams = params ?? getCurrentResolvedParams();
         const currentQueryKey = getQueryKey(effectiveParams);
@@ -509,12 +515,13 @@ export function createQueryStore<
 
         if (abortInterruptedFetches) abortActiveFetch();
 
-        if (!options?.force && !disableCache) {
-          const { errorInfo, lastFetchedAt } = get().queryCache[currentQueryKey] ?? {};
+        if (!options?.force) {
+          const { errorInfo, lastFetchedAt: cachedLastFetchedAt } = get().queryCache[currentQueryKey] ?? {};
           const errorRetriesExhausted = errorInfo && errorInfo.retryCount >= maxRetries;
+          const lastFetchedAt = cachedLastFetchedAt || (disableCache && lastFetchKey === currentQueryKey ? get().lastFetchedAt : null);
 
           if ((!errorInfo || errorRetriesExhausted) && lastFetchedAt && Date.now() - lastFetchedAt <= (options?.staleTime ?? staleTime)) {
-            return;
+            return null;
           }
         }
 
@@ -523,10 +530,14 @@ export function createQueryStore<
 
         const fetchOperation = async () => {
           try {
+            if (enableLogs) console.log('[🔄 Fetching 🔄] for queryKey: ', currentQueryKey, '::: params: ', effectiveParams);
             const rawResult = await (abortInterruptedFetches ? fetchWithAbortControl(effectiveParams) : fetcher(effectiveParams));
+            const lastFetchedAt = Date.now();
+            if (enableLogs) console.log('[✅ Fetch Successful ✅] for queryKey: ', currentQueryKey);
+
             let transformedData: TData;
             try {
-              transformedData = transform ? transform(rawResult) : (rawResult as TData);
+              transformedData = transform ? transform(rawResult, effectiveParams) : (rawResult as TData);
             } catch (transformError) {
               throw new RainbowError(`[createQueryStore: ${persistConfig?.storageKey || currentQueryKey}]: transform failed`, {
                 cause: transformError,
@@ -534,7 +545,6 @@ export function createQueryStore<
             }
 
             set(state => {
-              const lastFetchedAt = Date.now();
               let newState: S = {
                 ...state,
                 error: null,
@@ -543,6 +553,13 @@ export function createQueryStore<
               };
 
               if (!setData && !disableCache) {
+                if (enableLogs)
+                  console.log(
+                    '[💾 Setting Cache 💾] for queryKey: ',
+                    currentQueryKey,
+                    'Has previous data? ',
+                    !!newState.queryCache[currentQueryKey]?.data
+                  );
                 newState.queryCache = {
                   ...newState.queryCache,
                   [currentQueryKey]: {
@@ -552,6 +569,7 @@ export function createQueryStore<
                   },
                 };
               } else if (setData) {
+                if (enableLogs) console.log('[💾 Setting Data 💾] for queryKey: ', currentQueryKey);
                 setData({
                   data: transformedData,
                   params: effectiveParams,
@@ -588,11 +606,16 @@ export function createQueryStore<
                 );
               }
             }
+
+            return transformedData ?? null;
           } catch (error) {
-            if (error === ABORT_ERROR) return;
+            if (error === ABORT_ERROR) {
+              if (enableLogs) console.log('[❌ Fetch Aborted ❌] for queryKey: ', currentQueryKey);
+              return null;
+            }
 
             const typedError = error instanceof Error ? error : new Error(String(error));
-            const entry = get().queryCache[currentQueryKey];
+            const entry = disableCache ? undefined : get().queryCache[currentQueryKey];
             const currentRetryCount = entry?.errorInfo?.retryCount ?? 0;
 
             onError?.(typedError, currentRetryCount);
@@ -600,50 +623,57 @@ export function createQueryStore<
             if (currentRetryCount < maxRetries) {
               if (get().subscriptionCount > 0) {
                 const errorRetryDelay = typeof retryDelay === 'function' ? retryDelay(currentRetryCount, typedError) : retryDelay;
-
-                activeRefetchTimeout = setTimeout(() => {
-                  const { enabled, fetch, subscriptionCount } = get();
-                  if (enabled && subscriptionCount > 0) {
-                    fetch(params, { force: true });
-                  }
-                }, errorRetryDelay);
+                if (errorRetryDelay !== Infinity) {
+                  activeRefetchTimeout = setTimeout(() => {
+                    const { enabled, fetch, subscriptionCount } = get();
+                    if (enabled && subscriptionCount > 0) {
+                      fetch(params, { force: true });
+                    }
+                  }, errorRetryDelay);
+                }
               }
 
               set(state => ({
                 ...state,
                 error: typedError,
-                status: QueryStatuses.Error,
                 queryCache: {
                   ...state.queryCache,
                   [currentQueryKey]: {
                     ...entry,
-                    errorState: {
+                    errorInfo: {
                       error: typedError,
+                      lastFailedAt: Date.now(),
                       retryCount: currentRetryCount + 1,
                     },
                   },
                 },
+                status: QueryStatuses.Error,
               }));
             } else {
+              /* Max retries exhausted */
               set(state => ({
                 ...state,
-                status: QueryStatuses.Error,
+                error: typedError,
                 queryCache: {
                   ...state.queryCache,
                   [currentQueryKey]: {
                     ...entry,
-                    errorState: {
+                    errorInfo: {
                       error: typedError,
-                      retryCount: currentRetryCount,
+                      lastFailedAt: Date.now(),
+                      retryCount: maxRetries,
                     },
                   },
                 },
+                status: QueryStatuses.Error,
               }));
             }
 
             logger.error(new RainbowError(`[createQueryStore: ${persistConfig?.storageKey || currentQueryKey}]: Failed to fetch data`), {
               error: typedError,
             });
+
+            return null;
           } finally {
             activeFetchPromise = null;
             lastFetchKey = null;
@@ -663,7 +693,8 @@ export function createQueryStore<
       getStatus() {
         const { queryKey, status } = get();
         const lastFetchedAt =
-          get().queryCache[queryKey]?.lastFetchedAt || (disableCache && lastFetchKey === queryKey ? get().lastFetchedAt : null);
+          (!disableCache && get().queryCache[queryKey]?.lastFetchedAt) ||
+          (disableCache && lastFetchKey === queryKey ? get().lastFetchedAt : null);
 
         return {
           isError: status === QueryStatuses.Error,
@@ -677,7 +708,8 @@ export function createQueryStore<
       isDataExpired(cacheTimeOverride?: number) {
         const { queryKey } = get();
         const lastFetchedAt =
-          get().queryCache[queryKey]?.lastFetchedAt || (disableCache && lastFetchKey === queryKey ? get().lastFetchedAt : null);
+          (!disableCache && get().queryCache[queryKey]?.lastFetchedAt) ||
+          (disableCache && lastFetchKey === queryKey ? get().lastFetchedAt : null);
 
         if (!lastFetchedAt) return true;
         const effectiveCacheTime = cacheTimeOverride ?? cacheTime;
@@ -687,7 +719,8 @@ export function createQueryStore<
       isStale(staleTimeOverride?: number) {
         const { queryKey } = get();
         const lastFetchedAt =
-          get().queryCache[queryKey]?.lastFetchedAt || (disableCache && lastFetchKey === queryKey ? get().lastFetchedAt : null);
+          (!disableCache && get().queryCache[queryKey]?.lastFetchedAt) ||
+          (disableCache && lastFetchKey === queryKey ? get().lastFetchedAt : null);
 
         if (!lastFetchedAt) return true;
         const effectiveStaleTime = staleTimeOverride ?? staleTime;
@@ -818,15 +851,15 @@ export function createQueryStore<
     if (!attachVal) continue;
 
     const subscribeFn = attachValueSubscriptionMap.get(attachVal);
-    if (ENABLE_LOGS) console.log('[🌀 ParamSubscription 🌀] Subscribed to param:', k);
+    if (enableLogs) console.log('[🌀 Param Subscription 🌀] Subscribed to param:', k);
 
     if (subscribeFn) {
       let oldVal = attachVal.value;
       const unsub = subscribeFn(() => {
         const newVal = attachVal.value;
-        if (!Object.is(oldVal, newVal)) {
+        if (!equal(oldVal, newVal)) {
           oldVal = newVal;
-          if (ENABLE_LOGS) console.log('[🌀 ParamChange 🌀] Param changed:', k);
+          if (enableLogs) console.log('[🌀 Param Change 🌀] Param changed:', k, '::: Old value:', oldVal, '::: New value:', newVal);
           onParamChange();
         }
       });

From d44633f24e76c7ba72f93cca571820efa9321a56 Mon Sep 17 00:00:00 2001
From: Christian Baroni <7061887+christianbaroni@users.noreply.github.com>
Date: Wed, 1 Jan 2025 10:22:55 +0000
Subject: [PATCH 27/29] Expose abortController to fetcher, use fast-compare for
 the default signal equalityFn, bump fast-compare

---
 package.json                           |  2 +-
 src/state/internal/createQueryStore.ts |  6 +++---
 src/state/internal/signal.ts           |  5 +++--
 yarn.lock                              | 10 +++++-----
 4 files changed, 12 insertions(+), 11 deletions(-)

diff --git a/package.json b/package.json
index 016bfd20be7..eccb06349e9 100644
--- a/package.json
+++ b/package.json
@@ -211,7 +211,7 @@
     "querystring-es3": "0.2.1",
     "react": "18.2.0",
     "react-coin-icon": "rainbow-me/react-coin-icon#06464588a3d986f6ef3a7d7341b2d7ea0c5ac50b",
-    "react-fast-compare": "2.0.4",
+    "react-fast-compare": "3.2.2",
     "react-fast-memo": "2.0.1",
     "react-flatten-children": "1.1.2",
     "react-freeze": "1.0.4",
diff --git a/src/state/internal/createQueryStore.ts b/src/state/internal/createQueryStore.ts
index b2cbd119068..aaeaff80ce6 100644
--- a/src/state/internal/createQueryStore.ts
+++ b/src/state/internal/createQueryStore.ts
@@ -173,7 +173,7 @@ export type QueryStoreConfig<TQueryFnData, TParams extends Record<string, unknow
    * A function responsible for fetching data from a remote source.
    * Receives parameters of type `TParams` and returns either a promise or a raw data value of type `TQueryFnData`.
    */
-  fetcher: (params: TParams) => TQueryFnData | Promise<TQueryFnData>;
+  fetcher: (params: TParams, abortController: AbortController | null) => TQueryFnData | Promise<TQueryFnData>;
   /**
    * A callback invoked whenever a fetch operation fails.
    * Receives the error and the current retry count.
@@ -452,7 +452,7 @@ export function createQueryStore<
       return await new Promise((resolve, reject) => {
         abortController.signal.addEventListener('abort', () => reject(ABORT_ERROR), { once: true });
 
-        Promise.resolve(fetcher(params)).then(resolve, reject);
+        Promise.resolve(fetcher(params, abortController)).then(resolve, reject);
       });
     } finally {
       if (activeAbortController === abortController) {
@@ -531,7 +531,7 @@ export function createQueryStore<
         const fetchOperation = async () => {
           try {
             if (enableLogs) console.log('[🔄 Fetching 🔄] for queryKey: ', currentQueryKey, '::: params: ', effectiveParams);
-            const rawResult = await (abortInterruptedFetches ? fetchWithAbortControl(effectiveParams) : fetcher(effectiveParams));
+            const rawResult = await (abortInterruptedFetches ? fetchWithAbortControl(effectiveParams) : fetcher(effectiveParams, null));
             const lastFetchedAt = Date.now();
             if (enableLogs) console.log('[✅ Fetch Successful ✅] for queryKey: ', currentQueryKey);
 
diff --git a/src/state/internal/signal.ts b/src/state/internal/signal.ts
index 536f0ba5690..c3eec56f9fc 100644
--- a/src/state/internal/signal.ts
+++ b/src/state/internal/signal.ts
@@ -1,3 +1,4 @@
+import equal from 'react-fast-compare';
 import { StoreApi } from 'zustand';
 
 const ENABLE_LOGS = false;
@@ -30,7 +31,7 @@ export function $<T, S>(store: StoreApi<T>, selector: (state: T) => S, equalityF
 export function $(
   store: StoreApi<unknown>,
   selector: (state: unknown) => unknown = identity,
-  equalityFn: (a: unknown, b: unknown) => boolean = Object.is
+  equalityFn: (a: unknown, b: unknown) => boolean = equal
 ) {
   return getOrCreateAttachValue(store, selector, equalityFn);
 }
@@ -44,7 +45,7 @@ const updateValue = <T>(obj: T, path: unknown[], value: unknown): T => {
   const [first, ...rest] = path;
   const prevValue = (obj as Record<string, unknown>)[first as string];
   const nextValue = updateValue(prevValue, rest, value);
-  if (Object.is(prevValue, nextValue)) {
+  if (equal(prevValue, nextValue)) {
     return obj;
   }
   const copied = Array.isArray(obj) ? obj.slice() : { ...obj };
diff --git a/yarn.lock b/yarn.lock
index 0c47a4cd051..43be201e984 100644
--- a/yarn.lock
+++ b/yarn.lock
@@ -8868,7 +8868,7 @@ __metadata:
     querystring-es3: "npm:0.2.1"
     react: "npm:18.2.0"
     react-coin-icon: "rainbow-me/react-coin-icon#06464588a3d986f6ef3a7d7341b2d7ea0c5ac50b"
-    react-fast-compare: "npm:2.0.4"
+    react-fast-compare: "npm:3.2.2"
     react-fast-memo: "npm:2.0.1"
     react-flatten-children: "npm:1.1.2"
     react-freeze: "npm:1.0.4"
@@ -22105,10 +22105,10 @@ __metadata:
   languageName: node
   linkType: hard
 
-"react-fast-compare@npm:2.0.4":
-  version: 2.0.4
-  resolution: "react-fast-compare@npm:2.0.4"
-  checksum: 10c0/f0300c677e95198b5f993cbb8a983dab09586157dc678f9e2b5b29ff941b6677a8776fbbdc425ce102fad86937e36bb45cfcfd797f006270b97ccf287ebfb885
+"react-fast-compare@npm:3.2.2":
+  version: 3.2.2
+  resolution: "react-fast-compare@npm:3.2.2"
+  checksum: 10c0/0bbd2f3eb41ab2ff7380daaa55105db698d965c396df73e6874831dbafec8c4b5b08ba36ff09df01526caa3c61595247e3269558c284e37646241cba2b90a367
   languageName: node
   linkType: hard
 

From 13af587bf52d4e69039b1944e89abd3cd0ef8498 Mon Sep 17 00:00:00 2001
From: Christian Baroni <7061887+christianbaroni@users.noreply.github.com>
Date: Wed, 1 Jan 2025 23:28:15 +0000
Subject: [PATCH 28/29] Add skipStoreUpdates fetch option, only null out
 lastFetchKey following a failed fetch

---
 src/state/internal/createQueryStore.ts | 32 +++++++++++++++++++++-----
 1 file changed, 26 insertions(+), 6 deletions(-)

diff --git a/src/state/internal/createQueryStore.ts b/src/state/internal/createQueryStore.ts
index aaeaff80ce6..2cb1c8bb2ae 100644
--- a/src/state/internal/createQueryStore.ts
+++ b/src/state/internal/createQueryStore.ts
@@ -53,6 +53,13 @@ interface FetchOptions {
    * If `true`, the fetch operation bypasses existing cached data.
    */
   force?: boolean;
+  /**
+   * If `true`, the fetch will simply return the data without any internal handling or side effects,
+   * running in parallel with any other ongoing fetches. Use together with `force: true` if you want to
+   * guarantee that a fresh fetch is triggered regardless of the current store state.
+   * @default false
+   */
+  skipStoreUpdates?: boolean;
   /**
    * Overrides the default stale duration for this fetch, in milliseconds.
    * If the fetch is successful, the subsequently scheduled refetch will occur after
@@ -513,7 +520,7 @@ export function createQueryStore<
           return activeFetchPromise;
         }
 
-        if (abortInterruptedFetches) abortActiveFetch();
+        if (abortInterruptedFetches && !options?.skipStoreUpdates) abortActiveFetch();
 
         if (!options?.force) {
           const { errorInfo, lastFetchedAt: cachedLastFetchedAt } = get().queryCache[currentQueryKey] ?? {};
@@ -525,13 +532,17 @@ export function createQueryStore<
           }
         }
 
-        set(state => ({ ...state, error: null, status: QueryStatuses.Loading }));
-        lastFetchKey = currentQueryKey;
+        if (!options?.skipStoreUpdates) {
+          set(state => ({ ...state, error: null, status: QueryStatuses.Loading }));
+          lastFetchKey = currentQueryKey;
+        }
 
         const fetchOperation = async () => {
           try {
             if (enableLogs) console.log('[🔄 Fetching 🔄] for queryKey: ', currentQueryKey, '::: params: ', effectiveParams);
-            const rawResult = await (abortInterruptedFetches ? fetchWithAbortControl(effectiveParams) : fetcher(effectiveParams, null));
+            const rawResult = await (abortInterruptedFetches && !options?.skipStoreUpdates
+              ? fetchWithAbortControl(effectiveParams)
+              : fetcher(effectiveParams, null));
             const lastFetchedAt = Date.now();
             if (enableLogs) console.log('[✅ Fetch Successful ✅] for queryKey: ', currentQueryKey);
 
@@ -544,6 +555,8 @@ export function createQueryStore<
               });
             }
 
+            if (options?.skipStoreUpdates) return transformedData;
+
             set(state => {
               let newState: S = {
                 ...state,
@@ -614,10 +627,18 @@ export function createQueryStore<
               return null;
             }
 
+            if (options?.skipStoreUpdates) {
+              logger.error(new RainbowError(`[createQueryStore: ${persistConfig?.storageKey || currentQueryKey}]: Failed to fetch data`), {
+                error,
+              });
+              return null;
+            }
+
             const typedError = error instanceof Error ? error : new Error(String(error));
             const entry = disableCache ? undefined : get().queryCache[currentQueryKey];
             const currentRetryCount = entry?.errorInfo?.retryCount ?? 0;
 
+            if (lastFetchKey === currentQueryKey) lastFetchKey = null;
             onError?.(typedError, currentRetryCount);
 
             if (currentRetryCount < maxRetries) {
@@ -670,13 +691,12 @@ export function createQueryStore<
             }
 
             logger.error(new RainbowError(`[createQueryStore: ${persistConfig?.storageKey || currentQueryKey}]: Failed to fetch data`), {
-              error: typedError,
+              error,
             });
 
             return null;
           } finally {
             activeFetchPromise = null;
-            lastFetchKey = null;
           }
         };
 

From f5af6a34954b21d25c3e9e84ab6f2ce4c7ada00b Mon Sep 17 00:00:00 2001
From: Christian Baroni <7061887+christianbaroni@users.noreply.github.com>
Date: Fri, 3 Jan 2025 23:02:07 +0000
Subject: [PATCH 29/29] Defer initial fetch when params aren't yet initialized,
 prevent redundant work when multiple subscribed components are mounted,
 reduce subscription handler redundancy, expose more tools to onFetched, clean
 up lastFetchedAt logic, other small improvements

---
 src/state/internal/createQueryStore.ts | 106 ++++++++++++++++---------
 1 file changed, 70 insertions(+), 36 deletions(-)

diff --git a/src/state/internal/createQueryStore.ts b/src/state/internal/createQueryStore.ts
index 2cb1c8bb2ae..91032e9208a 100644
--- a/src/state/internal/createQueryStore.ts
+++ b/src/state/internal/createQueryStore.ts
@@ -190,7 +190,12 @@ export type QueryStoreConfig<TQueryFnData, TParams extends Record<string, unknow
    * A callback invoked whenever fresh data is successfully fetched.
    * Receives the transformed data and the store's set function, which can optionally be used to update store state.
    */
-  onFetched?: (data: TData, set: (partial: S | Partial<S> | ((state: S) => S | Partial<S>)) => void) => void;
+  onFetched?: (info: {
+    data: TData;
+    fetch: (params?: TParams, options?: FetchOptions) => Promise<TData | null>;
+    params: TParams;
+    set: (partial: S | Partial<S> | ((state: S) => S | Partial<S>)) => void;
+  }) => void;
   /**
    * A function that overrides the default behavior of setting the fetched data in the store's query cache.
    * Receives an object containing the transformed data, the query parameters, the query key, and the store's set function.
@@ -419,8 +424,10 @@ export function createQueryStore<
     );
   }
 
-  let directValues: Partial<TParams> = {};
-  let paramAttachVals: Partial<Record<keyof TParams, AttachValue<unknown>>> = {};
+  let directValues: Partial<TParams> | null = null;
+  let paramAttachVals: Partial<Record<keyof TParams, AttachValue<unknown>>> | null = null;
+  let paramUnsubscribes: Unsubscribe[] = [];
+  let fetchAfterParamCreation = false;
 
   let activeAbortController: AbortController | null = null;
   let activeFetchPromise: Promise<TData | null> | null = null;
@@ -442,7 +449,7 @@ export function createQueryStore<
   const getQueryKey = (params: TParams): string => JSON.stringify(Object.values(params));
 
   const getCurrentResolvedParams = () => {
-    const currentParams = { ...directValues };
+    const currentParams: Partial<TParams> = directValues ? { ...directValues } : {};
     for (const k in paramAttachVals) {
       const attachVal = paramAttachVals[k as keyof TParams];
       if (!attachVal) continue;
@@ -496,8 +503,7 @@ export function createQueryStore<
 
       const currentQueryKey = getQueryKey(params);
       const lastFetchedAt =
-        (!disableCache && get().queryCache[currentQueryKey]?.lastFetchedAt) ||
-        (disableCache && lastFetchKey === currentQueryKey ? get().lastFetchedAt : null);
+        (disableCache ? lastFetchKey === currentQueryKey && get().lastFetchedAt : get().queryCache[currentQueryKey]?.lastFetchedAt) || null;
       const timeUntilRefetch = lastFetchedAt ? effectiveStaleTime - (Date.now() - lastFetchedAt) : effectiveStaleTime;
 
       activeRefetchTimeout = setTimeout(() => {
@@ -510,11 +516,13 @@ export function createQueryStore<
 
     const baseMethods = {
       async fetch(params: TParams | undefined, options: FetchOptions | undefined) {
-        if (!options?.force && !get().enabled) return null;
+        const { enabled, status } = get();
+
+        if (!options?.force && !enabled) return null;
 
         const effectiveParams = params ?? getCurrentResolvedParams();
         const currentQueryKey = getQueryKey(effectiveParams);
-        const isLoading = get().status === QueryStatuses.Loading;
+        const isLoading = status === QueryStatuses.Loading;
 
         if (activeFetchPromise && lastFetchKey === currentQueryKey && isLoading && !options?.force) {
           return activeFetchPromise;
@@ -523,12 +531,24 @@ export function createQueryStore<
         if (abortInterruptedFetches && !options?.skipStoreUpdates) abortActiveFetch();
 
         if (!options?.force) {
-          const { errorInfo, lastFetchedAt: cachedLastFetchedAt } = get().queryCache[currentQueryKey] ?? {};
+          /* Check for valid cached data */
+          const {
+            lastFetchedAt: storeLastFetchedAt,
+            queryCache: { [currentQueryKey]: cacheEntry },
+            subscriptionCount,
+          } = get();
+
+          const { errorInfo, lastFetchedAt: cachedLastFetchedAt } = cacheEntry ?? {};
           const errorRetriesExhausted = errorInfo && errorInfo.retryCount >= maxRetries;
-          const lastFetchedAt = cachedLastFetchedAt || (disableCache && lastFetchKey === currentQueryKey ? get().lastFetchedAt : null);
+          const lastFetchedAt = (disableCache ? lastFetchKey === currentQueryKey && storeLastFetchedAt : cachedLastFetchedAt) || null;
+          const isStale = !lastFetchedAt || Date.now() - lastFetchedAt >= (options?.staleTime ?? staleTime);
 
-          if ((!errorInfo || errorRetriesExhausted) && lastFetchedAt && Date.now() - lastFetchedAt <= (options?.staleTime ?? staleTime)) {
-            return null;
+          if (!isStale && (!errorInfo || errorRetriesExhausted)) {
+            if (!activeRefetchTimeout && subscriptionCount > 0 && staleTime !== 0 && staleTime !== Infinity) {
+              scheduleNextFetch(effectiveParams, options);
+            }
+            if (enableLogs) console.log('[💾 Returning Cached Data 💾] for queryKey: ', currentQueryKey);
+            return cacheEntry?.data ?? null;
           }
         }
 
@@ -539,7 +559,7 @@ export function createQueryStore<
 
         const fetchOperation = async () => {
           try {
-            if (enableLogs) console.log('[🔄 Fetching 🔄] for queryKey: ', currentQueryKey, '::: params: ', effectiveParams);
+            if (enableLogs) console.log('[🔄 Fetching 🔄] for queryKey: ', currentQueryKey, ':: params: ', effectiveParams);
             const rawResult = await (abortInterruptedFetches && !options?.skipStoreUpdates
               ? fetchWithAbortControl(effectiveParams)
               : fetcher(effectiveParams, null));
@@ -555,7 +575,10 @@ export function createQueryStore<
               });
             }
 
-            if (options?.skipStoreUpdates) return transformedData;
+            if (options?.skipStoreUpdates) {
+              if (enableLogs) console.log('[🥷 Successful Parallel Fetch 🥷] for queryKey: ', currentQueryKey);
+              return transformedData;
+            }
 
             set(state => {
               let newState: S = {
@@ -610,7 +633,7 @@ export function createQueryStore<
 
             if (onFetched) {
               try {
-                onFetched(transformedData, set);
+                onFetched({ data: transformedData, fetch: baseMethods.fetch, params: effectiveParams, set });
               } catch (onFetchedError) {
                 logger.error(
                   new RainbowError(`[createQueryStore: ${persistConfig?.storageKey || currentQueryKey}]: onFetched callback failed`, {
@@ -713,8 +736,7 @@ export function createQueryStore<
       getStatus() {
         const { queryKey, status } = get();
         const lastFetchedAt =
-          (!disableCache && get().queryCache[queryKey]?.lastFetchedAt) ||
-          (disableCache && lastFetchKey === queryKey ? get().lastFetchedAt : null);
+          (disableCache ? lastFetchKey === queryKey && get().lastFetchedAt : get().queryCache[queryKey]?.lastFetchedAt) || null;
 
         return {
           isError: status === QueryStatuses.Error,
@@ -728,8 +750,7 @@ export function createQueryStore<
       isDataExpired(cacheTimeOverride?: number) {
         const { queryKey } = get();
         const lastFetchedAt =
-          (!disableCache && get().queryCache[queryKey]?.lastFetchedAt) ||
-          (disableCache && lastFetchKey === queryKey ? get().lastFetchedAt : null);
+          (disableCache ? lastFetchKey === queryKey && get().lastFetchedAt : get().queryCache[queryKey]?.lastFetchedAt) || null;
 
         if (!lastFetchedAt) return true;
         const effectiveCacheTime = cacheTimeOverride ?? cacheTime;
@@ -739,8 +760,7 @@ export function createQueryStore<
       isStale(staleTimeOverride?: number) {
         const { queryKey } = get();
         const lastFetchedAt =
-          (!disableCache && get().queryCache[queryKey]?.lastFetchedAt) ||
-          (disableCache && lastFetchKey === queryKey ? get().lastFetchedAt : null);
+          (disableCache ? lastFetchKey === queryKey && get().lastFetchedAt : get().queryCache[queryKey]?.lastFetchedAt) || null;
 
         if (!lastFetchedAt) return true;
         const effectiveStaleTime = staleTimeOverride ?? staleTime;
@@ -762,10 +782,12 @@ export function createQueryStore<
     const subscribeWithSelector = api.subscribe;
 
     api.subscribe = (listener: (state: S, prevState: S) => void) => {
-      set(state => ({ ...state, subscriptionCount: state.subscriptionCount + 1 }));
-      const unsubscribe = subscribeWithSelector(listener);
+      set(state => ({
+        ...state,
+        subscriptionCount: state.subscriptionCount + 1,
+      }));
 
-      const handleSetEnabled = subscribeWithSelector((state: S, prevState: S) => {
+      const handleSetEnabled = (state: S, prevState: S) => {
         if (state.enabled !== prevState.enabled) {
           if (state.enabled) {
             const currentParams = getCurrentResolvedParams();
@@ -784,22 +806,34 @@ export function createQueryStore<
             }
           }
         }
+      };
+
+      const unsubscribe = subscribeWithSelector((state: S, prevState: S) => {
+        listener(state, prevState);
+        handleSetEnabled(state, prevState);
       });
 
-      const { enabled, fetch, isStale } = get();
-      const currentParams = getCurrentResolvedParams();
-      const currentQueryKey = getQueryKey(currentParams);
+      const shouldWaitForParams = params !== undefined && !directValues && !paramAttachVals;
+
+      if (shouldWaitForParams) {
+        fetchAfterParamCreation = true;
+      } else {
+        const { enabled, fetch, isStale, subscriptionCount } = get();
+        const currentParams = getCurrentResolvedParams();
+        const currentQueryKey = getQueryKey(currentParams);
 
-      set(state => ({ ...state, queryKey: currentQueryKey }));
+        set(state => ({ ...state, queryKey: currentQueryKey }));
 
-      if (!get().queryCache[currentQueryKey] || isStale()) {
-        fetch(currentParams);
-      } else if (enabled) {
-        scheduleNextFetch(currentParams, undefined);
+        if (subscriptionCount === 1 && enabled) {
+          if (isStale() || !get().queryCache[currentQueryKey]?.lastFetchedAt) {
+            fetch(currentParams);
+          } else {
+            scheduleNextFetch(currentParams, undefined);
+          }
+        }
       }
 
       return () => {
-        handleSetEnabled();
         unsubscribe();
         set(state => {
           const newCount = Math.max(state.subscriptionCount - 1, 0);
@@ -864,8 +898,6 @@ export function createQueryStore<
     queryCapableStore.fetch();
   };
 
-  let paramUnsubscribes: Unsubscribe[] = [];
-
   for (const k in paramAttachVals) {
     const attachVal = paramAttachVals[k];
     if (!attachVal) continue;
@@ -879,7 +911,7 @@ export function createQueryStore<
         const newVal = attachVal.value;
         if (!equal(oldVal, newVal)) {
           oldVal = newVal;
-          if (enableLogs) console.log('[🌀 Param Change 🌀] Param changed:', k, '::: Old value:', oldVal, '::: New value:', newVal);
+          if (enableLogs) console.log('[🌀 Param Change 🌀] Param changed:', k, ':: Old value:', oldVal, ':: New value:', newVal);
           onParamChange();
         }
       });
@@ -887,6 +919,8 @@ export function createQueryStore<
     }
   }
 
+  if (fetchAfterParamCreation) onParamChange();
+
   return queryCapableStore;
 }