@rozenite/tanstack-query-plugin 1.7.0-rc.2 → 1.8.0

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);

package/tsconfig.json

@@ -11,12 +11,15 @@
     "noFallthroughCasesInSwitch": true,
     "module": "ESNext",
     "moduleResolution": "bundler",
+    "paths": {
+      "@rozenite/agent-shared": ["../agent-shared/src/index.ts"]
+    },
     "resolveJsonModule": true,
     "isolatedModules": true,
     "noEmit": true,
     "jsx": "react-jsx"
   },
-  "include": ["src/**/*", "react-native.ts", "rozenite.config.ts"],
+  "include": ["src/**/*", "react-native.ts", "sdk.ts", "rozenite.config.ts"],
   "exclude": ["node_modules", "dist", "build"],
   "references": [
     {

package/vite.config.ts

@@ -1,10 +1,17 @@
 /// <reference types='vitest' />
+import { resolve } from 'node:path';
 import { defineConfig } from 'vite';
 import { rozenitePlugin } from '@rozenite/vite-plugin';
 
 export default defineConfig({
   root: __dirname,
   plugins: [rozenitePlugin()],
+  test: {
+    passWithNoTests: true,
+    alias: {
+      '@rozenite/agent-shared': resolve(__dirname, '../agent-shared/src/index.ts'),
+    },
+  },
   base: './',
   build: {
     outDir: './dist',