@rozenite/tanstack-query-plugin 1.0.0-alpha.4 → 1.0.0-alpha.6

package/src/react-native/useSyncTanStackCache.ts

@@ -1,23 +1,236 @@
-import { useQueryClient } from '@tanstack/react-query';
-import { useEffect } from 'react';
+import {
+  QueryClient,
+  QueryCacheNotifyEvent,
+  MutationCacheNotifyEvent,
+} from '@tanstack/react-query';
+import { useEffect, useMemo, useRef } from 'react';
+import equal from 'fast-deep-equal';
 import { TanStackQueryPluginClient } from '../shared/messaging';
-import { dehydrateQueryClient } from '../shared/dehydrate';
+import {
+  dehydrateQuery,
+  dehydrateMutation,
+  dehydrateObservers,
+} from '../shared/dehydrate';
+import { SerializableObserver, PartialQueryState } from '../shared/types';
 
 export const useSyncTanStackCache = (
+  queryClient: QueryClient,
   client: TanStackQueryPluginClient | null
 ) => {
-  const queryClient = useQueryClient();
+  const previousObserversRef = useRef<Map<string, SerializableObserver[]>>(
+    new Map()
+  );
+
+  const handler = useMemo(() => {
+    return (event: QueryCacheNotifyEvent | MutationCacheNotifyEvent): void => {
+      if (!client) {
+        return;
+      }
+
+      if (event.type === 'observerResultsUpdated') {
+        // We don't need to sync this type of events.
+        return;
+      }
+
+      if ('query' in event) {
+        const { query, type } = event;
+
+        if (type === 'added' || type === 'removed') {
+          if (type === 'removed') {
+            // We don't need to store observers for removed queries.
+            previousObserversRef.current.delete(query.queryHash);
+          }
+
+          const dehydratedQuery = dehydrateQuery(query);
+          client.send('sync-query-event', { type, data: dehydratedQuery });
+          return;
+        }
+
+        if (type === 'updated' && 'action' in event) {
+          const action = event.action;
+
+          switch (action.type) {
+            case 'fetch': {
+              // Only need loading state and fetch metadata
+              const fetchData: PartialQueryState = {
+                queryHash: query.queryHash,
+                state: {
+                  status: query.state.status,
+                  fetchStatus: query.state.fetchStatus,
+                  fetchMeta: query.state.fetchMeta,
+                  dataUpdatedAt: query.state.dataUpdatedAt,
+                  errorUpdatedAt: query.state.errorUpdatedAt,
+                },
+              };
+              client.send('sync-query-event', {
+                type: 'updated',
+                action: 'fetch',
+                data: fetchData,
+              });
+              break;
+            }
+
+            case 'success': {
+              // Only need new data and success state
+              const successData: PartialQueryState = {
+                queryHash: query.queryHash,
+                state: {
+                  status: query.state.status,
+                  data: query.state.data,
+                  dataUpdatedAt: query.state.dataUpdatedAt,
+                  error: query.state.error,
+                  errorUpdatedAt: query.state.errorUpdatedAt,
+                  fetchStatus: query.state.fetchStatus,
+                },
+              };
+              client.send('sync-query-event', {
+                type: 'updated',
+                action: 'success',
+                data: successData,
+              });
+              break;
+            }
+
+            case 'error': {
+              // Only need error information
+              const errorData: PartialQueryState = {
+                queryHash: query.queryHash,
+                state: {
+                  status: query.state.status,
+                  error: query.state.error,
+                  errorUpdatedAt: query.state.errorUpdatedAt,
+                  fetchStatus: query.state.fetchStatus,
+                },
+              };
+              client.send('sync-query-event', {
+                type: 'updated',
+                action: 'error',
+                data: errorData,
+              });
+              break;
+            }
+
+            case 'setState': {
+              // Only need the specific state changes
+              const setStateData: PartialQueryState = {
+                queryHash: query.queryHash,
+                state: action.state, // Only the changed state parts
+              };
+              client.send('sync-query-event', {
+                type: 'updated',
+                action: 'setState',
+                data: setStateData,
+              });
+              break;
+            }
+
+            case 'invalidate': {
+              // Only need invalidation flag
+              const invalidateData: PartialQueryState = {
+                queryHash: query.queryHash,
+                state: {
+                  isInvalidated: query.state.isInvalidated,
+                },
+              };
+              client.send('sync-query-event', {
+                type: 'updated',
+                action: 'invalidate',
+                data: invalidateData,
+              });
+              break;
+            }
+
+            case 'pause': {
+              // Need to sync the paused fetch status
+              const pauseData: PartialQueryState = {
+                queryHash: query.queryHash,
+                state: {
+                  fetchStatus: query.state.fetchStatus,
+                },
+              };
+              client.send('sync-query-event', {
+                type: 'updated',
+                action: 'pause',
+                data: pauseData,
+              });
+              break;
+            }
+
+            case 'continue': {
+              // Need to sync the resumed fetch status
+              const continueData: PartialQueryState = {
+                queryHash: query.queryHash,
+                state: {
+                  fetchStatus: query.state.fetchStatus,
+                },
+              };
+              client.send('sync-query-event', {
+                type: 'updated',
+                action: 'continue',
+                data: continueData,
+              });
+              break;
+            }
+
+            default: {
+              // Fallback to full data for unknown actions
+              const dehydratedQuery = dehydrateQuery(query);
+              client.send('sync-query-event', { type, data: dehydratedQuery });
+            }
+          }
+          return;
+        }
+
+        if (
+          type === 'observerAdded' ||
+          type === 'observerRemoved' ||
+          type === 'observerOptionsUpdated'
+        ) {
+          const dehydratedObservers = dehydrateObservers(query);
+          const previousObservers = previousObserversRef.current.get(
+            query.queryHash
+          );
+
+          // For observerOptionsUpdated, only send if the observers actually changed
+          if (type === 'observerOptionsUpdated' && previousObservers) {
+            if (equal(previousObservers, dehydratedObservers)) {
+              // No meaningful change, skip sending
+              return;
+            }
+          }
+
+          previousObserversRef.current.set(
+            query.queryHash,
+            dehydratedObservers
+          );
+
+          client.send('sync-query-event', {
+            type,
+            data: {
+              queryHash: query.queryHash,
+              observers: dehydratedObservers,
+            },
+          });
+          return;
+        }
+      }
+
+      const { mutation, type } = event;
+      if (mutation) {
+        const dehydratedMutation = dehydrateMutation(mutation);
+        client.send('sync-mutation-event', {
+          type,
+          data: dehydratedMutation,
+        });
+      }
+    };
+  }, [client]);
 
   useEffect(() => {
     if (!client) {
       return;
     }
 
-    const handler = (): void => {
-      const dehydratedState = dehydrateQueryClient(queryClient);
-      client.send('sync-data', { data: dehydratedState });
-    };
-
     const mutationCacheSubscription = queryClient
       .getMutationCache()
       .subscribe(handler);
@@ -29,5 +242,5 @@ export const useSyncTanStackCache = (
       mutationCacheSubscription();
       queryCacheSubscription();
     };
-  }, [client, queryClient]);
+  }, [client, queryClient, handler]);
 };

package/src/react-native/{useTanStackQueryDevTools.ts → useTanstackQueryDevTools.ts}

@@ -1,3 +1,4 @@
+import { QueryClient } from '@tanstack/react-query';
 import { useRozeniteDevToolsClient } from '@rozenite/plugin-bridge';
 import { TanStackQueryPluginEventMap } from '../shared/messaging';
 import { useSyncOnlineStatus } from '../shared/useSyncOnlineStatus';
@@ -5,16 +6,16 @@ import { useHandleDevToolsMessages } from './useHandleDevToolsMessages';
 import { useSyncTanStackCache } from './useSyncTanStackCache';
 import { useHandleInitialData } from './useHandleInitialData';
 
-export const useTanStackQueryDevTools = () => {
+export const useTanStackQueryDevTools = (queryClient: QueryClient) => {
   const client = useRozeniteDevToolsClient<TanStackQueryPluginEventMap>({
     pluginId: '@rozenite/tanstack-query-plugin',
   });
 
   useSyncOnlineStatus(client);
 
-  useHandleDevToolsMessages(client);
+  useHandleDevToolsMessages(queryClient, client);
 
-  useSyncTanStackCache(client);
+  useSyncTanStackCache(queryClient, client);
 
-  useHandleInitialData(client);
+  useHandleInitialData(queryClient, client);
 };

package/src/shared/dehydrate.ts

@@ -6,13 +6,15 @@ import type {
   SerializableQueryClient,
 } from './types';
 
+export const dehydrateObservers = (query: Query): SerializableObserver[] => {
+  return query.observers.map((observer) => ({
+    queryHash: query.queryHash,
+    options: observer.options,
+  }));
+};
+
 export const dehydrateQuery = (query: Query): SerializableQuery => {
-  const dehydratedObservers: SerializableObserver[] = query.observers.map(
-    (observer) => ({
-      queryHash: query.queryHash,
-      options: observer.options,
-    })
-  );
+  const dehydratedObservers: SerializableObserver[] = dehydrateObservers(query);
 
   return {
     state: query.state,

package/src/shared/hydrate.ts

@@ -7,100 +7,244 @@ import {
   QueryObserverOptions,
   QueryState,
 } from '@tanstack/react-query';
-import { SerializableQueryClient } from './types';
+import type {
+  SerializableQuery,
+  SerializableMutation,
+  SerializableQueryClient,
+  SerializableObserver,
+  PartialQueryState,
+} from './types';
 
 const mockQueryFn = () => {
   return Promise.resolve(null);
 };
 
-export const hydrateQueryClient = (
+const hydrateObservers = (
   client: QueryClient,
-  dehydratedState: SerializableQueryClient
-): void => {
-  const queryCache = client.getQueryCache();
-  const mutationCache = client.getMutationCache();
+  queryHash: string,
+  dehydratedObservers: SerializableObserver[]
+) => {
+  const query = client.getQueryCache().get(queryHash);
+
+  if (!query) {
+    return;
+  }
+
+  query.observers.forEach((observer) => {
+    query.removeObserver(observer);
+  });
 
-  // Sync mutations
-  dehydratedState.mutations.forEach(({ options, state }) => {
-    const existingMutation = mutationCache.find({
-      mutationKey: options.mutationKey,
-    });
-    const hydratedState: MutationState<unknown, Error, void, unknown> =
-      state as MutationState<unknown, Error, void, unknown>;
-    const hydratedOptions: MutationOptions = options;
-
-    if (existingMutation) {
-      mutationCache.remove(existingMutation);
+  dehydratedObservers.forEach((observerState) => {
+    const hydratedOptions: InfiniteQueryObserverOptions | QueryObserverOptions =
+      observerState.options;
+
+    if ('initialPageParam' in hydratedOptions) {
+      delete hydratedOptions.initialPageParam;
+    }
+
+    if ('behavior' in hydratedOptions) {
+      delete hydratedOptions.behavior;
     }
 
-    mutationCache.build(client, hydratedOptions, hydratedState);
+    hydratedOptions.queryFn = mockQueryFn;
+
+    const observer = new QueryObserver(
+      client,
+      hydratedOptions as QueryObserverOptions
+    );
+    query.addObserver(observer);
+  });
+};
+
+const hydrateMutation = (
+  client: QueryClient,
+  dehydratedMutation: SerializableMutation
+) => {
+  const mutationCache = client.getMutationCache();
+  const { options, state } = dehydratedMutation;
+
+  const existingMutation = mutationCache.find({
+    mutationKey: options.mutationKey,
   });
+  const hydratedState: MutationState<unknown, Error, void, unknown> =
+    state as MutationState<unknown, Error, void, unknown>;
+  const hydratedOptions: MutationOptions = options;
+
+  if (existingMutation) {
+    mutationCache.remove(existingMutation);
+  }
+
+  mutationCache.build(client, hydratedOptions, hydratedState);
+};
+
+const hydrateQuery = (
+  client: QueryClient,
+  dehydratedQuery: SerializableQuery
+) => {
+  const queryCache = client.getQueryCache();
+  const { queryKey, state, queryHash, observers } = dehydratedQuery;
+
+  let query = queryCache.get(queryHash);
+  const data = state.data;
+  const hydratedState: QueryState = state;
+
+  // Do not hydrate if an existing query exists with newer data
+  if (query) {
+    if (
+      query.state.dataUpdatedAt < state.dataUpdatedAt ||
+      query.state.fetchStatus !== state.fetchStatus
+    ) {
+      query.setState({
+        ...hydratedState,
+        data,
+      });
+      query.setOptions({
+        ...query.options,
+        queryFn: mockQueryFn,
+        retry: 0,
+      });
+    }
+  } else {
+    // Restore query
+    query = queryCache.build(
+      client,
+      {
+        ...client.getDefaultOptions().hydrate?.queries,
+        queryKey,
+        queryHash,
+        queryFn: mockQueryFn,
+      },
+      {
+        ...hydratedState,
+        data,
+      }
+    );
+  }
+
+  hydrateObservers(client, queryHash, observers);
+};
+
+export const hydrateQueryClient = (
+  client: QueryClient,
+  dehydratedState: SerializableQueryClient
+): void => {
+  // Hydrate mutations
+  dehydratedState.mutations.forEach((mutation) =>
+    hydrateMutation(client, mutation)
+  );
 
   // Hydrate queries
-  dehydratedState.queries.forEach(
-    ({ queryKey, state, queryHash, observers }) => {
-      let query = queryCache.get(queryHash);
-      const data = state.data;
-      const hydratedState: QueryState = state;
+  dehydratedState.queries.forEach((query) => hydrateQuery(client, query));
+};
 
-      // Do not hydrate if an existing query exists with newer data
-      if (query) {
-        if (
-          query.state.dataUpdatedAt < state.dataUpdatedAt ||
-          query.state.fetchStatus !== state.fetchStatus
-        ) {
-          query.setState({
-            ...hydratedState,
-            data,
-          });
-          query.setOptions({
-            ...query.options,
-            queryFn: mockQueryFn,
-            retry: 0,
-          });
-        }
+export const applyQueryEvent = (
+  queryClient: QueryClient,
+  type:
+    | 'added'
+    | 'updated'
+    | 'removed'
+    | 'observerAdded'
+    | 'observerRemoved'
+    | 'observerResultsUpdated'
+    | 'observerOptionsUpdated',
+  data: SerializableQuery | PartialQueryState,
+  action?: string
+): void => {
+  const queryCache = queryClient.getQueryCache();
+
+  switch (type) {
+    case 'added': {
+      hydrateQuery(queryClient, data as SerializableQuery);
+      break;
+    }
+    case 'updated': {
+      if (action && 'state' in data && !('queryKey' in data)) {
+        // Handle action-based partial state update
+        applyPartialQueryState(queryClient, data as PartialQueryState);
       } else {
-        // Restore query
-        query = queryCache.build(
-          client,
-          {
-            ...client.getDefaultOptions().hydrate?.queries,
-            queryKey,
-            queryHash,
-            queryFn: mockQueryFn,
-          },
-          {
-            ...hydratedState,
-            data,
-          }
-        );
+        // Handle full query update (backward compatibility)
+        hydrateQuery(queryClient, data as SerializableQuery);
       }
+      break;
+    }
+    case 'observerAdded':
+    case 'observerRemoved':
+    case 'observerOptionsUpdated': {
+      hydrateObservers(queryClient, data.queryHash, (data as any).observers);
+      break;
+    }
+    case 'removed': {
+      const query = queryCache.get(data.queryHash);
+      if (query) {
+        queryCache.remove(query);
+      }
+      break;
+    }
+  }
+};
 
-      query.observers.forEach((observer) => {
-        query.removeObserver(observer);
-      });
-
-      observers.forEach((observerState) => {
-        const hydratedOptions:
-          | InfiniteQueryObserverOptions
-          | QueryObserverOptions = observerState.options;
+// New function to handle action-based partial state updates
+export const applyPartialQueryState = (
+  queryClient: QueryClient,
+  data: PartialQueryState
+): void => {
+  const queryCache = queryClient.getQueryCache();
+  const query = queryCache.get(data.queryHash);
 
-        if ('initialPageParam' in hydratedOptions) {
-          delete hydratedOptions.initialPageParam;
-        }
+  if (!query) {
+    // Query doesn't exist, we need to create it with minimal data
+    // This is a fallback for when we receive a partial update for a non-existent query
+    return;
+  }
 
-        if ('behavior' in hydratedOptions) {
-          delete hydratedOptions.behavior;
-        }
+  // Apply the partial state changes
+  if (data.state) {
+    query.setState(data.state);
+  }
+};
 
-        hydratedOptions.queryFn = mockQueryFn;
+export const applyMutationEvent = (
+  queryClient: QueryClient,
+  type:
+    | 'added'
+    | 'updated'
+    | 'removed'
+    | 'observerAdded'
+    | 'observerRemoved'
+    | 'observerResultsUpdated'
+    | 'observerOptionsUpdated',
+  data: SerializableMutation
+): void => {
+  const mutationCache = queryClient.getMutationCache();
 
-        const observer = new QueryObserver(
-          client,
-          hydratedOptions as QueryObserverOptions
-        );
-        query.addObserver(observer);
+  switch (type) {
+    case 'added':
+    case 'updated':
+    case 'observerAdded':
+    case 'observerRemoved':
+    case 'observerResultsUpdated':
+    case 'observerOptionsUpdated': {
+      hydrateMutation(queryClient, data);
+      break;
+    }
+    case 'removed': {
+      const mutation = mutationCache.find({
+        mutationKey: data.options.mutationKey,
       });
+      if (mutation) {
+        mutationCache.remove(mutation);
+      }
+      break;
     }
-  );
+  }
+};
+
+export const applyQueryObserverEvent = (
+  queryClient: QueryClient,
+  data: {
+    queryHash: string;
+    observers: SerializableObserver[];
+  }
+): void => {
+  hydrateObservers(queryClient, data.queryHash, data.observers);
 };

package/src/shared/messaging.ts

@@ -1,5 +1,12 @@
 import type { RozeniteDevToolsClient } from '@rozenite/plugin-bridge';
-import { DevToolsActionType, SerializableQueryClient } from './types';
+import {
+  DevToolsActionType,
+  SerializableQueryClient,
+  SerializableQuery,
+  SerializableMutation,
+  SerializableObserver,
+  PartialQueryState,
+} from './types';
 
 export type TanStackQueryPluginEventMap = {
   'online-status-changed': {
@@ -13,6 +20,34 @@ export type TanStackQueryPluginEventMap = {
   'sync-data': {
     data: SerializableQueryClient;
   };
+  'sync-query-event':
+    | {
+        type: 'added' | 'removed';
+        data: SerializableQuery;
+      }
+    | {
+        type: 'updated';
+        action?: string; // Add action type for optimization
+        data: SerializableQuery | PartialQueryState; // Support both full and partial state
+      }
+    | {
+        type: 'observerAdded' | 'observerRemoved' | 'observerOptionsUpdated';
+        data: {
+          queryHash: string;
+          observers: SerializableObserver[];
+        };
+      };
+  'sync-mutation-event': {
+    type:
+      | 'added'
+      | 'updated'
+      | 'removed'
+      | 'observerAdded'
+      | 'observerRemoved'
+      | 'observerResultsUpdated'
+      | 'observerOptionsUpdated';
+    data: SerializableMutation;
+  };
 };
 
 export type TanStackQueryPluginClient =

package/src/shared/types.ts

@@ -14,6 +14,11 @@ export type SerializableQuery = {
   observers: SerializableObserver[];
 };
 
+export type PartialQueryState = {
+  queryHash: string;
+  state: Partial<QueryState>;
+};
+
 export type SerializableMutation = {
   mutationId: number;
   options: MutationOptions;

package/src/shared/useSyncOnlineStatus.ts

@@ -11,14 +11,13 @@ export const useSyncOnlineStatus = (
     }
 
     const onlineManagerSubscription = onlineManager.subscribe((online) => {
-      // client.send('online-status-changed', { online });
+      client.send('online-status-changed', { online });
     });
 
     const onlineMessageSubscription = client.onMessage(
       'online-status-changed',
       ({ online }) => {
-        console.log('changed from ', onlineManager.isOnline(), 'to', online);
-        // onlineManager.setOnline(online);
+        onlineManager.setOnline(online);
       }
     );