diff --git a/apps/web/src/api/buster-electric/messages/hooks.ts b/apps/web/src/api/buster-electric/messages/hooks.ts index 97d96deb3..6837f566d 100644 --- a/apps/web/src/api/buster-electric/messages/hooks.ts +++ b/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 + ); +}; diff --git a/apps/web/src/api/buster-electric/messages/shapes.ts b/apps/web/src/api/buster-electric/messages/shapes.ts index 28d7197a0..5bc687886 100644 --- a/apps/web/src/api/buster-electric/messages/shapes.ts +++ b/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 => { return { - params: { table: 'messages', where: `chat_id='${chatId}'`, columns } + params: { + table: 'messages', + where: `chat_id='${chatId}'`, + columns: columns || MESSAGE_DEFAULT_COLUMNS + } }; }; diff --git a/apps/web/src/api/buster_rest/chats/queryRequests.ts b/apps/web/src/api/buster_rest/chats/queryRequests.ts index a8d5c3805..d859f8327 100644 --- a/apps/web/src/api/buster_rest/chats/queryRequests.ts +++ b/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[0], queryClientProp?: QueryClient ) => { @@ -170,6 +170,20 @@ export const prefetchGetChat = async ( return queryClient; }; +export const prefetchGetChat = async ( + params: Parameters[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 || {}; diff --git a/apps/web/src/layouts/ChatLayout/ChatContext/useChatStreaming.tsx b/apps/web/src/layouts/ChatLayout/ChatContext/useChatStreaming.tsx index f326ff384..8326d5a0f 100644 --- a/apps/web/src/layouts/ChatLayout/ChatContext/useChatStreaming.tsx +++ b/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,