feat: tanstack query plugin improvements, refetch control & state

jest.config.json

@@ -3,6 +3,7 @@
     "testEnvironment": "node",
     "modulePathIgnorePatterns": ["<rootDir>/dist"],
     "moduleNameMapper": {
+        "@legendapp/state/sync-plugins/tanstack-query": "<rootDir>/src/sync-plugins/tanstack-query",
         "@legendapp/state/sync-plugins/crud": "<rootDir>/src/sync-plugins/crud",
         "@legendapp/state/sync": "<rootDir>/sync",
         "@legendapp/state/config/configureLegendState": "<rootDir>/src/config/configureLegendState",

src/sync-plugins/tanstack-query.ts

@@ -9,6 +9,7 @@ import {
     QueryKey,
     QueryObserver,
     QueryObserverOptions,
+    QueryObserverResult,
     notifyManager,
 } from '@tanstack/query-core';
 
@@ -19,11 +20,20 @@ export interface ObservableQueryOptions<TQueryFnData, TError, TData, TQueryKey e
     queryKey?: TQueryKey | (() => TQueryKey);
 }
 
+export interface QueryState<TError = DefaultError> {
+    isLoading: boolean;
+    isFetching: boolean;
+    error: TError | null;
+    status: 'pending' | 'error' | 'success';
+    fetchStatus: 'fetching' | 'paused' | 'idle';
+}
+
 export interface SyncedQueryParams<TQueryFnData, TError, TData, TQueryKey extends QueryKey>
     extends Omit<SyncedOptions<TData>, 'get' | 'set' | 'retry'> {
     queryClient: QueryClient;
     query: ObservableQueryOptions<TQueryFnData, TError, TData, TQueryKey>;
     mutation?: MutationObserverOptions<TQueryFnData, TError, TData>;
+    onQueryStateChange?: (state: QueryState<TError>) => void;
 }
 
 export function syncedQuery<
@@ -32,7 +42,14 @@ export function syncedQuery<
     TData = TQueryFnData,
     TQueryKey extends QueryKey = QueryKey,
 >(params: SyncedQueryParams<TQueryFnData, TError, TData, TQueryKey>): Synced<TData> {
-    const { query: options, mutation: mutationOptions, queryClient, initial: initialParam, ...rest } = params;
+    const {
+        query: options,
+        mutation: mutationOptions,
+        queryClient,
+        initial: initialParam,
+        onQueryStateChange,
+        ...rest
+    } = params;
 
     if (initialParam !== undefined) {
         const initialValue = isFunction(initialParam) ? initialParam() : initialParam;
@@ -84,35 +101,79 @@ export function syncedQuery<
     observer = new Observer!<TQueryFnData, TError, TData, TQueryKey>(queryClient!, latestOptions);
 
     let isFirstRun = true;
+    let rejectInitialPromise: undefined | ((error: Error) => void) = undefined;
+    // Track whether subscribe was just called in this sync cycle.
+    // The sync infrastructure calls subscribe() before get() on (re-)observation,
+    // but skips subscribe() on explicit sync() when already subscribed.
+    let subscribedInThisCycle = false;
+
     const get = (async () => {
+        const wasJustSubscribed = subscribedInThisCycle;
+        subscribedInThisCycle = false;
+
         if (isFirstRun) {
             isFirstRun = false;
 
             // Get the initial optimistic results if it's already cached
             const result = observer!.getOptimisticResult(latestOptions);
 
             if (result.isLoading) {
-                await new Promise((resolve) => {
+                return await new Promise<TData>((resolve, reject) => {
                     resolveInitialPromise = resolve;
+                    rejectInitialPromise = reject;
                 });
             }
 
             return result.data!;
+        } else if (wasJustSubscribed) {
+            // Re-observation (remount): return cached data, let TQ observer
+            // handle refetch decisions via subscription (refetchOnMount, staleTime, etc.)
+            return observer!.getCurrentResult().data!;
         } else {
+            // Explicit sync(): always force refetch from the server
             return Promise.resolve(observer!.refetch()).then((res) => (res as any).data as TData);
         }
     }) as () => Promise<TData>;
 
-    const subscribe = ({ update }: SyncedSubscribeParams<TData>) => {
+    const emitQueryState = (result: QueryObserverResult<TData, TError>) => {
+        if (onQueryStateChange) {
+            onQueryStateChange({
+                isLoading: result.isLoading,
+                isFetching: result.isFetching,
+                error: result.error,
+                status: result.status,
+                fetchStatus: result.fetchStatus,
+            });
+        }
+    };
+
+    const subscribe = ({ update, onError, node }: SyncedSubscribeParams<TData>) => {
+        subscribedInThisCycle = true;
+
         // Subscribe to Query's observer and update the observable
         const unsubscribe = observer!.subscribe(
-            notifyManager.batchCalls((result) => {
+            notifyManager.batchCalls((result: QueryObserverResult<TData, TError>) => {
+                emitQueryState(result);
+
                 if (result.status === 'success') {
+                    // Clear error on success
+                    if (node.state && node.state.error.peek()) {
+                        node.state.error.set(undefined);
+                    }
                     if (resolveInitialPromise) {
                         resolveInitialPromise(result.data);
                         resolveInitialPromise = undefined;
+                        rejectInitialPromise = undefined;
                     }
                     update({ value: result.data });
+                } else if (result.status === 'error') {
+                    // Propagate error to syncState via onError
+                    if (rejectInitialPromise) {
+                        rejectInitialPromise(result.error as Error);
+                        rejectInitialPromise = undefined;
+                        resolveInitialPromise = undefined;
+                    }
+                    onError(result.error as Error);
                 }
             }),
         );

tests/sync-plugins/tanstack-query.mock.ts

@@ -0,0 +1,93 @@
+export function createQueryCoreMock() {
+    const refetchMock = jest.fn(() => Promise.resolve({ data: 'fresh', status: 'success' }));
+    let subscriberCallback: ((result: any) => void) | undefined;
+
+    const defaultResult = {
+        data: 'initial',
+        status: 'success',
+        isLoading: false,
+        isFetching: false,
+        isStale: false,
+        error: null,
+        fetchStatus: 'idle' as const,
+    };
+
+    class QueryObserver {
+        client: any;
+        options: any;
+        notifyOptions: any;
+
+        constructor(client: any, options: any) {
+            this.client = client;
+            this.options = options;
+        }
+
+        getOptimisticResult() {
+            return defaultResult;
+        }
+
+        getCurrentResult() {
+            return defaultResult;
+        }
+
+        setOptions(_options: any, _notifyOptions?: any) {
+            this.options = _options;
+            this.notifyOptions = _notifyOptions;
+        }
+
+        refetch() {
+            return refetchMock();
+        }
+
+        subscribe(cb: any) {
+            subscriberCallback = cb;
+            return () => {
+                subscriberCallback = undefined;
+            };
+        }
+
+        updateResult() {}
+    }
+
+    class MutationObserver {
+        client: any;
+        options: any;
+
+        constructor(client: any, options: any) {
+            this.client = client;
+            this.options = options;
+        }
+
+        mutate(value: any) {
+            return Promise.resolve(value);
+        }
+    }
+
+    class QueryClient {
+        defaultQueryOptions(options: any) {
+            return options;
+        }
+
+        getMutationCache() {
+            return { findAll: () => [], remove: () => {} };
+        }
+    }
+
+    const notifyManager = {
+        batchCalls:
+            (fn: (...args: any[]) => any) =>
+            (...args: any[]) =>
+                fn(...args),
+    };
+
+    return {
+        __esModule: true,
+        QueryObserver,
+        MutationObserver,
+        QueryClient,
+        notifyManager,
+        DefaultError: Error,
+        refetchMock,
+        getSubscriberCallback: () => subscriberCallback,
+    };
+}