From 4c1c68db9226dc676fc5261ad908459a98a6f70c Mon Sep 17 00:00:00 2001
From: Nate Kelley <nate@buster.so>
Date: Thu, 13 Feb 2025 11:18:45 -0700
Subject: [PATCH] Update useSocketQueryOn.tsx

---
 .../useSocketQueryNew/useSocketQueryOn.tsx    | 79 +++++++++++--------
 1 file changed, 46 insertions(+), 33 deletions(-)

diff --git a/web/src/hooks/useSocketQueryNew/useSocketQueryOn.tsx b/web/src/hooks/useSocketQueryNew/useSocketQueryOn.tsx
index ce7d0ac47..d39a798ca 100644
--- a/web/src/hooks/useSocketQueryNew/useSocketQueryOn.tsx
+++ b/web/src/hooks/useSocketQueryNew/useSocketQueryOn.tsx
@@ -7,22 +7,16 @@ import {
   UseQueryOptions,
   UseQueryResult
 } from '@tanstack/react-query';
-import type {
-  BusterSocketRequest,
-  BusterSocketResponse,
-  BusterSocketResponseRoute
-} from '@/api/buster_socket';
+import type { BusterSocketResponse, BusterSocketResponseRoute } from '@/api/buster_socket';
 import { useBusterWebSocket } from '@/context/BusterWebSocket';
-import { useEffect } from 'react';
-import {
-  BusterSocketResponseConfig,
-  BusterSocketResponseConfigRoute,
-  InferBusterSocketResponseData
-} from './types';
-import { useMount } from 'ahooks';
+import { useRef, useTransition } from 'react';
+import { InferBusterSocketResponseData } from './types';
+import { useMemoizedFn, useMount, useUnmount } from 'ahooks';
 import { queryOptionsConfig } from './queryKeyConfig';
 import { BusterChat } from '@/api/asset_interfaces';
 
+type UseSocketQueryOnResult<TData, TError> = UseQueryResult<TData, TError>;
+
 export const useSocketQueryOn = <
   TRoute extends BusterSocketResponseRoute,
   TError = unknown,
@@ -31,30 +25,58 @@ export const useSocketQueryOn = <
 >(
   socketResponse: TRoute,
   options: UseQueryOptions<TData, TError, TData, TQueryKey>,
-  callback?: (d: InferBusterSocketResponseData<TRoute>) => TData
-): UseQueryResult<TData, TError> => {
+  // callback?: (d: InferBusterSocketResponseData<TRoute>) => TData,
+  callback?: (currentData: TData | null, newData: InferBusterSocketResponseData<TRoute>) => TData
+): UseSocketQueryOnResult<TData, TError> => {
   const busterSocket = useBusterWebSocket();
   const queryClient = useQueryClient();
   const queryKey = options.queryKey;
+  const bufferRef = useRef<TData | null>(null);
+  const [isPending, startTransition] = useTransition();
+
+  const hasBufferCallback = !!callback;
+
+  const socketCallback = useMemoizedFn((d: unknown) => {
+    const socketData = d as InferBusterSocketResponseData<TRoute>;
+
+    const transformer = callback || defaultCallback<TData, TRoute>;
+    const transformedData = transformer(bufferRef.current, socketData);
+
+    if (hasBufferCallback) {
+      bufferRef.current = transformedData;
+      startTransition(() => {
+        queryClient.setQueryData<TData>(queryKey, transformedData);
+      });
+    } else {
+      queryClient.setQueryData<TData>(queryKey, transformedData);
+    }
+  });
 
   useMount(() => {
     busterSocket.on({
       route: socketResponse,
-      callback: (d: unknown) => {
-        const socketData = d as InferBusterSocketResponseData<TRoute>;
-        const transformer = callback || ((d: InferBusterSocketResponseData<TRoute>) => d as TData);
-        const transformedData = transformer(socketData);
-        queryClient.setQueryData<TData>(queryKey, transformedData);
-      }
+      callback: socketCallback
     } as BusterSocketResponse);
   });
 
+  useUnmount(() => {
+    busterSocket.off({
+      route: socketResponse,
+      callback: socketCallback
+    });
+  });
+
   return useQuery({
     ...options,
-    enabled: false
+    enabled: false // must be disabled, because it will be enabled by the socket
   });
 };
 
+const defaultCallback = <TData, TRoute extends BusterSocketResponseRoute>(
+  currentData: TData | null,
+  d: InferBusterSocketResponseData<TRoute>
+) => d as TData;
+
 const ExampleComponent = () => {
   const options = queryOptionsConfig['/chats/get:getChat']('123');
   const { data } = useSocketQueryOn('/chats/get:getChat', options);
@@ -70,15 +92,6 @@ const ExampleComponent = () => {
 };
 
 // Create fresh options for delete chat that match the expected BusterChat type
-const deleteChatInitialData: BusterChat = {
-  id: '123',
-  title: '',
-  is_favorited: false,
-  messages: [],
-  created_at: new Date().toISOString(),
-  updated_at: new Date().toISOString(),
-  created_by: '',
-  created_by_id: '',
-  created_by_name: '',
-  created_by_avatar: null
-};
+const deleteChatInitialData = {
+  id: '123'
+} as unknown as BusterChat;