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/23] 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/23] 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/23] 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/23] 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/23] 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/23] 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/23] 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/23] 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/23] 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/23] 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/23] 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/23] 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/23] 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/23] 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/23] [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/23] 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/23] [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/23] 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/23] 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/23] 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/23] 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/23] 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/23] 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 {};