Merge pull request #562 from buster-so/devin/BUS-1447-1752896644

feat(BUS-1447): Add hook to track new message insertions from external sources

apps/web/src/api/buster-electric/messages/hooks.ts

@@ -4,7 +4,7 @@ import { useShape, useShapeStream } from '../instances';
 import { useChatUpdate } from '@/context/Chats/useChatUpdate';
 import { updateMessageShapeToIChatMessage } from './helpers';
 import { useMemoizedFn } from '@/hooks';
-import { prefetchGetChatsList, useGetChatMemoized } from '@/api/buster_rest/chats';
+import { useGetChatMemoized, useGetChatMessageMemoized } from '@/api/buster_rest/chats';
 import uniq from 'lodash/uniq';
 import type { ChatMessageResponseMessage_File } from '@buster/server-shared/chats';
 import type { BusterChatMessage } from '../../asset_interfaces/chat';
@@ -12,6 +12,7 @@ import { useQueryClient } from '@tanstack/react-query';
 import { dashboardQueryKeys } from '../../query_keys/dashboard';
 import isEmpty from 'lodash/isEmpty';
 import { metricsQueryKeys } from '../../query_keys/metric';
+import { chatQueryKeys } from '../../query_keys/chat';
 
 export const useGetMessage = ({ chatId, messageId }: { chatId: string; messageId: string }) => {
   const shape = useMemo(() => messageShape({ chatId, messageId }), [chatId, messageId]);
@@ -24,6 +25,7 @@ export const useGetMessages = ({ chatId }: { chatId: string }) => {
 };
 
 const updateOperations: Array<'insert' | 'update' | 'delete'> = ['update'];
+const insertOperations: Array<'insert' | 'update' | 'delete'> = ['insert'];
 
 export const useTrackAndUpdateMessageChanges = (
   {
@@ -40,6 +42,7 @@ export const useTrackAndUpdateMessageChanges = (
   const { onUpdateChatMessage, onUpdateChat } = useChatUpdate();
   const checkIfWeHaveAFollowupDashboard = useCheckIfWeHaveAFollowupDashboard(messageId);
   const getChatMemoized = useGetChatMemoized();
+  const queryClient = useQueryClient();
 
   const subscribe = !!chatId && !!messageId && messageId !== 'undefined';
 
@@ -63,6 +66,7 @@ export const useTrackAndUpdateMessageChanges = (
           if (currentMessageIds.length !== allMessageIds.length) {
             onUpdateChat({
               ...chat,
+              id: chatId,
               message_ids: allMessageIds
             });
           }
@@ -75,16 +79,20 @@ export const useTrackAndUpdateMessageChanges = (
               (reasoningMessage as ChatMessageResponseMessage_File)?.file_type === 'dashboard'
             );
           });
-          if (hasFiles) {
-            prefetchGetChatsList();
-          }
 
           if (!isEmpty(iChatMessage.response_message_ids)) {
             checkIfWeHaveAFollowupDashboard(iChatMessage);
           }
 
           if (iChatMessage.is_completed) {
-            prefetchGetChatsList();
+            queryClient.invalidateQueries({
+              queryKey: chatQueryKeys.chatsGetList().queryKey
+            });
+            if (hasFiles) {
+              queryClient.invalidateQueries({
+                queryKey: metricsQueryKeys.metricsGetList().queryKey
+              });
+            }
           }
         }
         callback?.(iChatMessage);
@@ -128,3 +136,46 @@ const useCheckIfWeHaveAFollowupDashboard = (messageId: string) => {
 
   return useMemoizedFn(method);
 };
+
+export const useTrackAndUpdateNewMessages = ({ chatId }: { chatId: string | undefined }) => {
+  const { onUpdateChat } = useChatUpdate();
+  const getChatMemoized = useGetChatMemoized();
+  const getChatMessageMemoized = useGetChatMessageMemoized();
+  const queryClient = useQueryClient();
+
+  const subscribe = !!chatId;
+
+  const shape = useMemo(() => messagesShape({ chatId: chatId || '', columns: ['id'] }), [chatId]);
+
+  return useShapeStream(
+    shape,
+    insertOperations,
+    useMemoizedFn((message) => {
+      if (message && message.value && chatId) {
+        const messageId = message.value.id;
+        const chat = getChatMemoized(chatId);
+
+        if (chat && messageId) {
+          const currentMessageIds = chat.message_ids;
+          const allMessageIds = uniq([...currentMessageIds, messageId]);
+
+          if (currentMessageIds.length !== allMessageIds.length) {
+            onUpdateChat({
+              ...chat,
+              id: chatId,
+              message_ids: allMessageIds
+            });
+
+            const messageIsStored = getChatMessageMemoized(messageId);
+            if (!messageIsStored) {
+              queryClient.invalidateQueries({
+                queryKey: chatQueryKeys.chatsGetChat(chatId).queryKey
+              });
+            }
+          }
+        }
+      }
+    }),
+    subscribe
+  );
+};

apps/web/src/api/buster-electric/messages/shapes.ts

@@ -15,7 +15,7 @@ export type BusterChatMessageShape = {
   is_completed: boolean;
 };
 
-const columns: (keyof BusterChatMessageShape)[] = [
+const MESSAGE_DEFAULT_COLUMNS: (keyof BusterChatMessageShape)[] = [
   'id',
   'response_messages',
   'reasoning',
@@ -36,18 +36,24 @@ export const messageShape = ({
     params: {
       table: 'messages',
       where: `chat_id='${chatId}' AND id='${messageId}'`,
-      columns,
+      columns: MESSAGE_DEFAULT_COLUMNS,
       replica: 'default'
     }
   };
 };
 
 export const messagesShape = ({
-  chatId
+  chatId,
+  columns = MESSAGE_DEFAULT_COLUMNS
 }: {
   chatId: string;
+  columns?: (keyof BusterChatMessageShape)[];
 }): ElectricShapeOptions<BusterChatMessageShape> => {
   return {
-    params: { table: 'messages', where: `chat_id='${chatId}'`, columns }
+    params: {
+      table: 'messages',
+      where: `chat_id='${chatId}'`,
+      columns: columns || MESSAGE_DEFAULT_COLUMNS
+    }
   };
 };

apps/web/src/api/buster_rest/chats/queryRequests.ts

@@ -152,7 +152,7 @@ export const useStartChatFromAsset = () => {
   });
 };
 
-export const prefetchGetChat = async (
+export const prefetchGetChatServer = async (
   params: Parameters<typeof getChat>[0],
   queryClientProp?: QueryClient
 ) => {
@@ -170,6 +170,20 @@ export const prefetchGetChat = async (
   return queryClient;
 };
 
+export const prefetchGetChat = async (
+  params: Parameters<typeof getChat>[0],
+  queryClientProp?: QueryClient
+) => {
+  const queryClient = queryClientProp || new QueryClient();
+
+  await queryClient.prefetchQuery({
+    ...chatQueryKeys.chatsGetChat(params.id),
+    queryFn: () => getChat(params)
+  });
+
+  return queryClient;
+};
+
 export const useUpdateChat = (params?: { updateToServer?: boolean }) => {
   const queryClient = useQueryClient();
   const { updateToServer = true } = params || {};

apps/web/src/layouts/ChatLayout/ChatContext/useChatStreaming.tsx

@@ -5,7 +5,7 @@ import { updateChatToIChat } from '@/lib/chat';
 import { useQueryClient } from '@tanstack/react-query';
 import { prefetchGetMetricDataClient } from '@/api/buster_rest/metrics';
 import { queryKeys } from '@/api/query_keys';
-import { useTrackAndUpdateMessageChanges } from '@/api/buster-electric/messages';
+import { useTrackAndUpdateMessageChanges, useTrackAndUpdateNewMessages } from '@/api/buster-electric/messages';
 import { useTrackAndUpdateChatChanges } from '@/api/buster-electric/chats';
 import { useEffect } from 'react';
 import { useGetChatMessageMemoized } from '@/api/buster_rest/chats';
@@ -90,6 +90,7 @@ export const useChatStreaming = ({
 
   //HOOKS FOR TRACKING CHAT AND MESSAGE CHANGES
   useTrackAndUpdateChatChanges({ chatId, isStreamingMessage });
+  useTrackAndUpdateNewMessages({ chatId });
   useTrackAndUpdateMessageChanges({ chatId, messageId, isStreamingMessage }, (c) => {
     const {
       reasoning_messages,