@rozenite/tanstack-query-plugin 1.7.0-rc.1 → 1.7.0

package/src/react-native/agent/__tests__/tanstack-query-agent.test.ts

@@ -5,9 +5,7 @@ import {
   QueryObserver,
   onlineManager,
 } from '@tanstack/react-query';
-import {
-  applyTanStackQueryDevtoolsAction,
-} from '../../devtools-actions';
+import { applyTanStackQueryDevtoolsAction } from '../../devtools-actions';
 import {
   createTanStackQueryAgentController,
   serializeForAgent,
@@ -30,7 +28,7 @@ const createQueryClient = () =>
 
 const subscribeController = (
   queryClient: QueryClient,
-  controller: ReturnType<typeof createTanStackQueryAgentController>
+  controller: ReturnType<typeof createTanStackQueryAgentController>,
 ) => {
   const unsubscribeQuery = queryClient
     .getQueryCache()
@@ -58,7 +56,7 @@ const createMutation = (
     context?: unknown;
     failureCount?: number;
     failureReason?: Error | null;
-  }
+  },
 ) => {
   return queryClient.getMutationCache().build(
     queryClient,
@@ -78,7 +76,7 @@ const createMutation = (
       status: options?.status ?? 'success',
       submittedAt: options?.submittedAt ?? Date.now(),
       variables: options?.variables,
-    } as any
+    } as any,
   ) as Mutation<unknown, Error, unknown, unknown>;
 };
 
@@ -126,8 +124,10 @@ describe('tanstack query agent controller', () => {
       controller.listQueries({
         limit: 1,
         cursor: firstPage.page.nextCursor,
-      })
-    ).toThrow('Cursor does not match the requested listing. Run the command again.');
+      }),
+    ).toThrow(
+      'Cursor does not match the requested listing. Run the command again.',
+    );
 
     unsubscribe();
   });
@@ -147,8 +147,10 @@ describe('tanstack query agent controller', () => {
       controller.listQueries({
         limit: 1,
         cursor: firstPage.page.nextCursor,
-      })
-    ).toThrow('Cursor does not match the requested listing. Run the command again.');
+      }),
+    ).toThrow(
+      'Cursor does not match the requested listing. Run the command again.',
+    );
 
     queryClient.setQueryData(['third'], { value: 3 });
     queryClient.setQueryData(['fourth'], { value: 4 });
@@ -159,8 +161,10 @@ describe('tanstack query agent controller', () => {
       controller.listQueries({
         limit: 1,
         cursor: secondPageSeed.page.nextCursor,
-      })
-    ).toThrow('Cursor does not match the requested listing. Run the command again.');
+      }),
+    ).toThrow(
+      'Cursor does not match the requested listing. Run the command again.',
+    );
 
     unsubscribe();
   });
@@ -194,8 +198,10 @@ describe('tanstack query agent controller', () => {
       controller.listMutations({
         limit: 1,
         cursor: firstPage.page.nextCursor,
-      })
-    ).toThrow('Cursor does not match the requested listing. Run the command again.');
+      }),
+    ).toThrow(
+      'Cursor does not match the requested listing. Run the command again.',
+    );
 
     const mutation = queryClient.getMutationCache().getAll()[0];
     queryClient.getMutationCache().remove(mutation);
@@ -207,8 +213,10 @@ describe('tanstack query agent controller', () => {
       controller.listMutations({
         limit: 1,
         cursor: secondPageSeed.page.nextCursor,
-      })
-    ).toThrow('Cursor does not match the requested listing. Run the command again.');
+      }),
+    ).toThrow(
+      'Cursor does not match the requested listing. Run the command again.',
+    );
 
     unsubscribe();
   });
@@ -234,8 +242,9 @@ describe('tanstack query agent controller', () => {
       queryFn,
     });
 
-    const queryHash = queryClient.getQueryCache().find({ queryKey: ['todos', 1] })!
-      .queryHash;
+    const queryHash = queryClient
+      .getQueryCache()
+      .find({ queryKey: ['todos', 1] })!.queryHash;
     const details = controller.getQueryDetails({ queryHash });
 
     expect(details.query).toMatchObject({
@@ -299,9 +308,9 @@ describe('tanstack query agent controller', () => {
       },
     });
 
-    expect(() =>
-      controller.getMutationDetails({ mutationId: 9999 })
-    ).toThrow('Unknown mutationId "9999"');
+    expect(() => controller.getMutationDetails({ mutationId: 9999 })).toThrow(
+      'Unknown mutationId "9999"',
+    );
 
     unsubscribe();
   });
@@ -310,9 +319,9 @@ describe('tanstack query agent controller', () => {
     const queryClient = createQueryClient();
     const controller = createTanStackQueryAgentController(queryClient);
 
-    expect(() =>
-      controller.getQueryDetails({ queryHash: 'missing' })
-    ).toThrow('Unknown queryHash "missing"');
+    expect(() => controller.getQueryDetails({ queryHash: 'missing' })).toThrow(
+      'Unknown queryHash "missing"',
+    );
   });
 
   it('gets and sets the online manager status', () => {
@@ -461,6 +470,37 @@ describe('tanstack query agent actions', () => {
     });
   });
 
+  it('sets query data through the shared UI dispatcher', async () => {
+    const queryClient = createQueryClient();
+
+    await queryClient.fetchQuery({
+      queryKey: ['editable-query'],
+      queryFn: async () => ({ value: 'initial' }),
+      initialData: { value: 'initial' },
+    });
+
+    const query = queryClient.getQueryCache().find({
+      queryKey: ['editable-query'],
+    })!;
+
+    const result = await applyTanStackQueryDevtoolsAction(queryClient, {
+      type: 'SET_QUERY_DATA',
+      queryHash: query.queryHash,
+      metadata: {
+        data: { value: 'changed' },
+      },
+    });
+
+    expect(result).toMatchObject({
+      applied: true,
+      action: 'SET_QUERY_DATA',
+      queryHash: query.queryHash,
+    });
+    expect(queryClient.getQueryData(['editable-query'])).toEqual({
+      value: 'changed',
+    });
+  });
+
   it('clears mutation cache through the shared UI/agent dispatcher', async () => {
     const queryClient = createQueryClient();
     createMutation(queryClient, {
@@ -501,8 +541,6 @@ describe('serializeForAgent', () => {
       ok: true,
       self: '[circular]',
     });
-    expect(serializeForAgent(new Example())).toBe(
-      '[non-serializable:Example]'
-    );
+    expect(serializeForAgent(new Example())).toBe('[non-serializable:Example]');
   });
 });

package/src/react-native/devtools-actions.ts

@@ -7,6 +7,9 @@ type QueryScopedActionInput = {
     'CLEAR_MUTATION_CACHE' | 'CLEAR_QUERY_CACHE'
   >;
   queryHash: string;
+  metadata?: {
+    data?: unknown;
+  };
 };
 
 type CacheScopedActionInput = {
@@ -33,7 +36,7 @@ const getActiveQuery = (queryClient: QueryClient, queryHash?: string) => {
 
 export const applyTanStackQueryDevtoolsAction = async (
   queryClient: QueryClient,
-  input: TanStackQueryDevtoolsActionInput
+  input: TanStackQueryDevtoolsActionInput,
 ) => {
   switch (input.type) {
     case 'CLEAR_QUERY_CACHE': {
@@ -49,8 +52,9 @@ export const applyTanStackQueryDevtoolsAction = async (
     }
 
     case 'CLEAR_MUTATION_CACHE': {
-      const mutationCountBefore =
-        queryClient.getMutationCache().getAll().length;
+      const mutationCountBefore = queryClient
+        .getMutationCache()
+        .getAll().length;
       queryClient.getMutationCache().clear();
       return {
         applied: true,
@@ -61,6 +65,16 @@ export const applyTanStackQueryDevtoolsAction = async (
       };
     }
 
+    case 'SET_QUERY_DATA': {
+      const activeQuery = getActiveQuery(queryClient, input.queryHash);
+      queryClient.setQueryData(activeQuery.queryKey, input.metadata?.data);
+      return {
+        applied: true,
+        action: input.type,
+        queryHash: activeQuery.queryHash,
+      };
+    }
+
     case 'TRIGGER_ERROR': {
       const activeQuery = getActiveQuery(queryClient, input.queryHash);
       const previousQueryOptions = activeQuery.options;

package/src/react-native/useHandleDevToolsMessages.ts

@@ -5,7 +5,7 @@ import { applyTanStackQueryDevtoolsAction } from './devtools-actions';
 
 export const useHandleDevToolsMessages = (
   queryClient: QueryClient,
-  client: TanStackQueryPluginClient | null
+  client: TanStackQueryPluginClient | null,
 ) => {
   useEffect(() => {
     if (!client) {
@@ -14,16 +14,19 @@ export const useHandleDevToolsMessages = (
 
     const subscription = client.onMessage(
       'devtools-action',
-      ({ type, queryHash }) => {
-        void applyTanStackQueryDevtoolsAction(queryClient, { type, queryHash })
-          .catch((error) => {
-            const message =
-              error instanceof Error ? error.message : String(error);
-            console.warn(
-              `[Rozenite, tanstack-query-plugin] Failed to apply devtools action "${type}": ${message}`
-            );
-          });
-      }
+      ({ type, queryHash, metadata }) => {
+        void applyTanStackQueryDevtoolsAction(queryClient, {
+          type,
+          queryHash,
+          metadata,
+        }).catch((error) => {
+          const message =
+            error instanceof Error ? error.message : String(error);
+          console.warn(
+            `[Rozenite, tanstack-query-plugin] Failed to apply devtools action "${type}": ${message}`,
+          );
+        });
+      },
     );
 
     return () => {

package/src/shared/hydrate.ts

@@ -14,6 +14,7 @@ import type {
   SerializableObserver,
   PartialQueryState,
 } from './types';
+import { applyRemoteQueryState, instrumentQuery } from './query-data-sync';
 
 const mockQueryFn = () => {
   return Promise.resolve(null);
@@ -22,7 +23,7 @@ const mockQueryFn = () => {
 const hydrateObservers = (
   client: QueryClient,
   queryHash: string,
-  dehydratedObservers: SerializableObserver[]
+  dehydratedObservers: SerializableObserver[],
 ) => {
   const query = client.getQueryCache().get(queryHash);
 
@@ -50,7 +51,7 @@ const hydrateObservers = (
 
     const observer = new QueryObserver(
       client,
-      hydratedOptions as QueryObserverOptions
+      hydratedOptions as QueryObserverOptions,
     );
     query.addObserver(observer);
   });
@@ -58,7 +59,7 @@ const hydrateObservers = (
 
 const hydrateMutation = (
   client: QueryClient,
-  dehydratedMutation: SerializableMutation
+  dehydratedMutation: SerializableMutation,
 ) => {
   const mutationCache = client.getMutationCache();
   const { options, state } = dehydratedMutation;
@@ -79,7 +80,7 @@ const hydrateMutation = (
 
 const hydrateQuery = (
   client: QueryClient,
-  dehydratedQuery: SerializableQuery
+  dehydratedQuery: SerializableQuery,
 ) => {
   const queryCache = client.getQueryCache();
   const { queryKey, state, queryHash, observers } = dehydratedQuery;
@@ -94,7 +95,7 @@ const hydrateQuery = (
       query.state.dataUpdatedAt < state.dataUpdatedAt ||
       query.state.fetchStatus !== state.fetchStatus
     ) {
-      query.setState({
+      applyRemoteQueryState(query, {
         ...hydratedState,
         data,
       });
@@ -117,20 +118,22 @@ const hydrateQuery = (
       {
         ...hydratedState,
         data,
-      }
+      },
     );
   }
 
+  instrumentQuery(query);
+
   hydrateObservers(client, queryHash, observers);
 };
 
 export const hydrateQueryClient = (
   client: QueryClient,
-  dehydratedState: SerializableQueryClient
+  dehydratedState: SerializableQueryClient,
 ): void => {
   // Hydrate mutations
   dehydratedState.mutations.forEach((mutation) =>
-    hydrateMutation(client, mutation)
+    hydrateMutation(client, mutation),
   );
 
   // Hydrate queries
@@ -148,7 +151,7 @@ export const applyQueryEvent = (
     | 'observerResultsUpdated'
     | 'observerOptionsUpdated',
   data: SerializableQuery | PartialQueryState,
-  action?: string
+  action?: string,
 ): void => {
   const queryCache = queryClient.getQueryCache();
 
@@ -186,7 +189,7 @@ export const applyQueryEvent = (
 // New function to handle action-based partial state updates
 export const applyPartialQueryState = (
   queryClient: QueryClient,
-  data: PartialQueryState
+  data: PartialQueryState,
 ): void => {
   const queryCache = queryClient.getQueryCache();
   const query = queryCache.get(data.queryHash);
@@ -199,7 +202,7 @@ export const applyPartialQueryState = (
 
   // Apply the partial state changes
   if (data.state) {
-    query.setState(data.state);
+    applyRemoteQueryState(query, data.state);
   }
 };
 
@@ -213,7 +216,7 @@ export const applyMutationEvent = (
     | 'observerRemoved'
     | 'observerResultsUpdated'
     | 'observerOptionsUpdated',
-  data: SerializableMutation
+  data: SerializableMutation,
 ): void => {
   const mutationCache = queryClient.getMutationCache();
 
@@ -244,7 +247,7 @@ export const applyQueryObserverEvent = (
   data: {
     queryHash: string;
     observers: SerializableObserver[];
-  }
+  },
 ): void => {
   hydrateObservers(queryClient, data.queryHash, data.observers);
 };

package/src/shared/messaging.ts

@@ -15,6 +15,9 @@ export type TanStackQueryPluginEventMap = {
   'devtools-action': {
     type: DevToolsActionType;
     queryHash: string;
+    metadata?: {
+      data?: unknown;
+    };
   };
   'request-initial-data': unknown;
   'sync-data': {

package/src/shared/query-data-sync.test.ts

@@ -0,0 +1,93 @@
+import { describe, expect, it, vi } from 'vitest';
+import { QueryClient } from '@tanstack/react-query';
+import {
+  applyRemoteQueryState,
+  instrumentQuery,
+  instrumentQueryClient,
+} from './query-data-sync';
+
+const createQueryClient = () =>
+  new QueryClient({
+    defaultOptions: {
+      queries: {
+        retry: false,
+        gcTime: Infinity,
+      },
+    },
+  });
+
+describe('query data sync instrumentation', () => {
+  it('emits query data updates for data-only setState edits', () => {
+    const queryClient = createQueryClient();
+    const handler = vi.fn();
+
+    queryClient.setQueryData(['edited'], { value: 1 });
+    const query = queryClient.getQueryCache().find({ queryKey: ['edited'] })!;
+
+    instrumentQuery(query, handler);
+
+    query.setState({
+      ...query.state,
+      data: { value: 2 },
+    });
+
+    expect(handler).toHaveBeenCalledTimes(1);
+    expect(handler).toHaveBeenCalledWith(
+      expect.objectContaining({
+        data: { value: 2 },
+      }),
+    );
+  });
+
+  it('emits query data updates for setQueryData user edits', () => {
+    const queryClient = createQueryClient();
+    const handler = vi.fn();
+
+    instrumentQueryClient(queryClient, handler);
+
+    queryClient.setQueryData(['edited'], { value: 1 });
+
+    expect(handler).toHaveBeenCalledTimes(1);
+    expect(handler).toHaveBeenCalledWith(
+      expect.objectContaining({
+        data: { value: 1 },
+      }),
+    );
+  });
+
+  it('does not emit query data updates for non-data state changes', () => {
+    const queryClient = createQueryClient();
+    const handler = vi.fn();
+
+    queryClient.setQueryData(['edited'], { value: 1 });
+    const query = queryClient.getQueryCache().find({ queryKey: ['edited'] })!;
+
+    instrumentQuery(query, handler);
+
+    query.setState({
+      data: undefined,
+      status: 'pending',
+      fetchStatus: 'fetching',
+    });
+
+    expect(handler).not.toHaveBeenCalled();
+  });
+
+  it('does not emit query data updates for remote state application', () => {
+    const queryClient = createQueryClient();
+    const handler = vi.fn();
+
+    queryClient.setQueryData(['edited'], { value: 1 });
+    const query = queryClient.getQueryCache().find({ queryKey: ['edited'] })!;
+
+    instrumentQuery(query, handler);
+
+    applyRemoteQueryState(query, {
+      ...query.state,
+      data: { value: 2 },
+    });
+
+    expect(handler).not.toHaveBeenCalled();
+    expect(query.state.data).toEqual({ value: 2 });
+  });
+});

package/src/shared/query-data-sync.ts

@@ -0,0 +1,156 @@
+import {
+  Query,
+  QueryClient,
+  QueryKey,
+  QueryState,
+} from '@tanstack/react-query';
+
+type QuerySetState = Query['setState'];
+type QueryClientSetQueryData = QueryClient['setQueryData'];
+
+type SyncQueryDataPayload = {
+  queryHash: string;
+  data: unknown;
+};
+
+type QueryDataSyncHandler = (payload: SyncQueryDataPayload) => void;
+
+const originalQuerySetStateMap = new WeakMap<Query, QuerySetState>();
+const originalSetQueryDataMap = new WeakMap<
+  QueryClient,
+  QueryClientSetQueryData
+>();
+const instrumentedQueryClients = new WeakSet<QueryClient>();
+const instrumentedQueries = new WeakSet<Query>();
+const queryClientHandlers = new WeakMap<
+  QueryClient,
+  QueryDataSyncHandler | undefined
+>();
+const queryHandlers = new WeakMap<Query, QueryDataSyncHandler | undefined>();
+
+const getOriginalQuerySetState = (query: Query): QuerySetState => {
+  return originalQuerySetStateMap.get(query) ?? query.setState.bind(query);
+};
+
+const emitIfDataChanged = (
+  query: Query,
+  previousData: unknown,
+  handler?: QueryDataSyncHandler,
+) => {
+  if (!handler || Object.is(previousData, query.state.data)) {
+    return;
+  }
+
+  handler({
+    queryHash: query.queryHash,
+    data: query.state.data,
+  });
+};
+
+const isDataOnlyStateUpdate = (
+  previousState: QueryState,
+  nextState: Partial<QueryState>,
+) => {
+  if (!('data' in nextState)) {
+    return false;
+  }
+
+  const changedKeys = Object.keys(nextState).filter((key) => {
+    const typedKey = key as keyof QueryState;
+    return !Object.is(previousState[typedKey], nextState[typedKey]);
+  }) as Array<keyof QueryState>;
+
+  return changedKeys.every((key) => key === 'data');
+};
+
+export const applyRemoteQueryState = (
+  query: Query,
+  state: Partial<QueryState>,
+) => {
+  const originalSetState = getOriginalQuerySetState(query);
+  originalSetState(state);
+};
+
+export const instrumentQuery = (
+  query: Query,
+  handler?: QueryDataSyncHandler,
+) => {
+  queryHandlers.set(query, handler);
+
+  if (instrumentedQueries.has(query)) {
+    return query;
+  }
+
+  const originalSetState = query.setState.bind(query);
+  originalQuerySetStateMap.set(query, originalSetState);
+
+  query.setState = ((state) => {
+    const shouldEmit = isDataOnlyStateUpdate(query.state, state);
+    const previousData = query.state.data;
+    const result = originalSetState(state);
+
+    if (shouldEmit) {
+      emitIfDataChanged(query, previousData, queryHandlers.get(query));
+    }
+
+    return result;
+  }) as QuerySetState;
+
+  instrumentedQueries.add(query);
+  return query;
+};
+
+const resolveQueryHash = (queryClient: QueryClient, queryKey: QueryKey) => {
+  return queryClient.getQueryCache().find({ queryKey })?.queryHash;
+};
+
+export const instrumentQueryClient = (
+  queryClient: QueryClient,
+  handler?: QueryDataSyncHandler,
+) => {
+  queryClientHandlers.set(queryClient, handler);
+
+  if (instrumentedQueryClients.has(queryClient)) {
+    queryClient
+      .getQueryCache()
+      .getAll()
+      .forEach((query) => {
+        instrumentQuery(query, handler);
+      });
+    return queryClient;
+  }
+
+  const originalSetQueryData = queryClient.setQueryData.bind(queryClient);
+  originalSetQueryDataMap.set(queryClient, originalSetQueryData);
+
+  queryClient.setQueryData = ((queryKey, updater, options) => {
+    const result = originalSetQueryData(queryKey, updater, options);
+    const query = queryClient.getQueryCache().find({ queryKey });
+    const activeQueryClientHandler = queryClientHandlers.get(queryClient);
+
+    if (query) {
+      instrumentQuery(query, activeQueryClientHandler);
+      const queryHash = resolveQueryHash(queryClient, queryKey);
+      const activeHandler = queryHandlers.get(query);
+
+      if (activeHandler && queryHash) {
+        activeHandler({
+          queryHash,
+          data: query.state.data,
+        });
+      }
+    }
+
+    return result;
+  }) as QueryClientSetQueryData;
+
+  queryClient
+    .getQueryCache()
+    .getAll()
+    .forEach((query) => {
+      instrumentQuery(query, handler);
+    });
+
+  instrumentedQueryClients.add(queryClient);
+  return queryClient;
+};

package/src/shared/types.ts

@@ -37,6 +37,7 @@ export type SerializableQueryClient = {
 
 export type DevToolsActionType =
   | 'REFETCH'
+  | 'SET_QUERY_DATA'
   | 'INVALIDATE'
   | 'RESET'
   | 'REMOVE'

package/src/ui/tanstack-query.tsx

@@ -1,3 +1,4 @@
+import { useEffect } from 'react';
 import { QueryClient, QueryClientProvider } from '@tanstack/react-query';
 import { ReactQueryDevtoolsPanel } from '@tanstack/react-query-devtools/production';
 import { useRozeniteDevToolsClient } from '@rozenite/plugin-bridge';
@@ -6,12 +7,23 @@ import { useSyncInitialData } from './useSyncInitialData';
 import { useSyncDevToolsEvents } from './useSyncDevToolsEvents';
 import { useSyncOnlineStatus } from '../shared/useSyncOnlineStatus';
 import { useHandleSyncMessages } from './useHandleSyncMessages';
+import { instrumentQueryClient } from '../shared/query-data-sync';
 
 const App = () => {
   const client = useRozeniteDevToolsClient<TanStackQueryPluginEventMap>({
     pluginId: '@rozenite/tanstack-query-plugin',
   });
 
+  useEffect(() => {
+    instrumentQueryClient(queryClient, ({ queryHash, data }) => {
+      client?.send('devtools-action', {
+        type: 'SET_QUERY_DATA',
+        queryHash,
+        metadata: { data },
+      });
+    });
+  }, [client]);
+
   useSyncInitialData(client);
 
   useSyncDevToolsEvents(client);