move chats to use new socket on syntax

web/src/api/buster_socket/chats/chatResponses.ts

@@ -15,7 +15,8 @@ export enum ChatsResponses {
   '/chats/post:generatingResponseMessage' = '/chats/post:generatingResponseMessage',
   '/chats/post:generatingReasoningMessage' = '/chats/post:generatingReasoningMessage',
   '/chats/post:complete' = '/chats/post:complete',
-  '/chats/delete:deleteChat' = '/chats/delete:deleteChat'
+  '/chats/delete:deleteChat' = '/chats/delete:deleteChat',
+  '/chats/update:updateChat' = '/chats/update:updateChat'
 }
 
 export type ChatList_getChatsList = {
@@ -49,6 +50,12 @@ export type Chat_deleteChat = {
   onError?: (d: unknown | RustApiError) => void;
 };
 
+export type Chat_updateChat = {
+  route: '/chats/update:updateChat';
+  callback: (d: BusterChat) => void;
+  onError?: (d: unknown | RustApiError) => void;
+};
+
 /***** CHAT PROGRESS EVENTS START ******/
 
 export type ChatPost_initializeChat = {

web/src/api/buster_socket_query/useSocketQueryEmitOn.tsx

@@ -9,7 +9,6 @@ import {
 import type { InferBusterSocketResponseData } from './types';
 import { useBusterWebSocket } from '@/context/BusterWebSocket';
 import { useEffect } from 'react';
-import { useMemoizedFn } from 'ahooks';
 import { useSocketQueryOn } from './useSocketQueryOn';
 import { timeout } from '@/utils';
 
@@ -41,7 +40,7 @@ export const useSocketQueryEmitOn = <
     }
   });
 
-  const queryResult = useSocketQueryOn(socketResponse, options, callback);
+  const queryResult = useSocketQueryOn({ socketResponse, options, callback });
 
   useEffect(() => {
     const queryState = queryClient.getQueryState(options.queryKey);

web/src/api/buster_socket_query/useSocketQueryOn.tsx

@@ -20,18 +20,23 @@ export const useSocketQueryOn = <
   TError = unknown,
   TData = InferBusterSocketResponseData<TRoute>,
   TQueryKey extends QueryKey = QueryKey
->(
-  socketResponse: TRoute,
-  options: UseQueryOptions<TData, TError, TData, TQueryKey>,
+>({
+  socketResponse,
+  options,
+  callback
+}: {
+  socketResponse: TRoute;
+  options?: UseQueryOptions<TData, TError, TData, TQueryKey> | null;
   callback?:
-    | ((currentData: TData | null, newData: InferBusterSocketResponseData<TRoute>) => TData)
-    | null
-): UseSocketQueryOnResult<TData, TError> => {
+    | ((
+        currentData: TData | null,
+        newData: InferBusterSocketResponseData<TRoute>
+      ) => TData | undefined | void)
+    | null;
+}): UseSocketQueryOnResult<TData, TError> => {
   const busterSocket = useBusterWebSocket();
   const queryClient = useQueryClient();
-  const queryKey = options.queryKey;
-  const bufferRef = useRef<TData | null>(null);
-  const [isPending, startTransition] = useTransition();
+  const queryKey = options?.queryKey ?? '';
 
   const hasBufferCallback = !!callback;
 
@@ -39,15 +44,13 @@ export const useSocketQueryOn = <
     const socketData = d as InferBusterSocketResponseData<TRoute>;
 
     const transformer = callback || defaultCallback<TData, TRoute>;
-    const transformedData = transformer(bufferRef.current, socketData);
+    const currentData = queryKey ? queryClient.getQueryData<TData>(queryKey)! : (null as TData);
+    const transformedData = transformer(currentData, socketData);
 
-    if (hasBufferCallback) {
-      bufferRef.current = transformedData;
-      startTransition(() => {
-        queryClient.setQueryData<TData>(queryKey, transformedData);
-      });
-    } else {
-      queryClient.setQueryData<TData>(queryKey, transformedData);
+    if (hasBufferCallback && transformedData) {
+      if (queryKey) queryClient.setQueryData<TData>(queryKey, transformedData);
+    } else if (transformedData) {
+      if (queryKey) queryClient.setQueryData<TData>(queryKey, transformedData);
     }
   });
 
@@ -69,6 +72,7 @@ export const useSocketQueryOn = <
 
   return useQuery({
     ...options,
+    queryKey: queryKey as TQueryKey,
     refetchOnMount: false,
     refetchOnWindowFocus: false,
     enabled: false // must be disabled, because it will be enabled by the socket

web/src/app/app/_controllers/ReasoningController/ReasoningController.tsx

@@ -2,7 +2,7 @@
 import React from 'react';
 import { useChatIndividualContextSelector } from '../../_layouts/ChatLayout/ChatContext';
 import { ReasoningMessageContainer } from './ReasoningMessageContainer';
-import { useBusterChatContextSelector } from '@/context/Chats';
+import { useMessageIndividual } from '@/context/Chats';
 
 interface ReasoningControllerProps {
   chatId: string;
@@ -11,7 +11,7 @@ interface ReasoningControllerProps {
 
 export const ReasoningController: React.FC<ReasoningControllerProps> = ({ chatId, messageId }) => {
   const hasChat = useChatIndividualContextSelector((state) => state.hasChat);
-  const message = useBusterChatContextSelector((state) => state.chatsMessages[messageId]);
+  const message = useMessageIndividual(messageId);
 
   if (!hasChat || !message) return null;
 

web/src/app/app/_layouts/ChatLayout/ChatContainer/ChatContent/ChatMessageBlock.tsx

@@ -1,16 +1,17 @@
 import React from 'react';
 import { ChatUserMessage } from './ChatUserMessage';
 import { ChatResponseMessages } from './ChatResponseMessages';
-import { useBusterChatContextSelector } from '@/context/Chats';
+import { useMessageIndividual } from '@/context/Chats';
 
 export const ChatMessageBlock: React.FC<{
   messageId: string;
 }> = React.memo(({ messageId }) => {
-  const message = useBusterChatContextSelector((state) => state.chatsMessages[messageId]);
-  const { request_message, response_messages, id, isCompletedStream, reasoning } = message || {};
+  const message = useMessageIndividual(messageId);
 
   if (!message) return null;
 
+  const { request_message, response_messages, id, isCompletedStream, reasoning } = message;
+
   return (
     <div className={'flex flex-col space-y-3.5 py-2 pl-4 pr-3'} id={id}>
       <ChatUserMessage requestMessage={request_message} />

web/src/app/app/_layouts/ChatLayout/ChatContext/ChatContext.tsx

@@ -5,9 +5,9 @@ import {
   useContextSelector
 } from '@fluentui/react-context-selector';
 import type { SelectedFile } from '../interfaces';
-import { useSubscribeIndividualChat } from './useSubscribeIndividualChat';
 import { useAutoChangeLayout } from './useAutoChangeLayout';
-import { useBusterChatContextSelector } from '@/context/Chats';
+import { useMessageIndividual } from '@/context/Chats';
+import { useSubscribeIndividualChat } from './useSubscribeIndividualChat';
 
 export const useChatIndividualContext = ({
   chatId,
@@ -36,9 +36,8 @@ export const useChatIndividualContext = ({
 
   //MESSAGES
   const currentMessageId = chatMessageIds[chatMessageIds.length - 1];
-  const isLoading = useBusterChatContextSelector(
-    (x) => !x.chatsMessages[currentMessageId]?.isCompletedStream
-  );
+  const message = useMessageIndividual(currentMessageId);
+  const isLoading = !message?.isCompletedStream;
 
   useAutoChangeLayout({
     lastMessageId: currentMessageId,

web/src/app/app/_layouts/ChatLayout/ChatContext/useAutoChangeLayout.ts

@@ -1,4 +1,4 @@
-import { useBusterChatContextSelector } from '@/context/Chats';
+import { useMessageIndividual } from '@/context/Chats';
 import type { SelectedFile } from '../interfaces';
 import { usePrevious } from 'ahooks';
 import { useEffect } from 'react';
@@ -11,7 +11,7 @@ export const useAutoChangeLayout = ({
   lastMessageId: string;
   onSetSelectedFile: (file: SelectedFile) => void;
 }) => {
-  const message = useBusterChatContextSelector((x) => x.chatsMessages[lastMessageId]);
+  const message = useMessageIndividual(lastMessageId);
   const reasoningMessagesLength = message?.reasoning?.length;
   const previousReasoningMessagesLength = usePrevious(reasoningMessagesLength);
   const isCompletedStream = message?.isCompletedStream;

web/src/app/app/_layouts/ChatLayout/ChatContext/useSubscribeIndividualChat.ts

@@ -1,8 +1,6 @@
 import type { IBusterChat } from '@/context/Chats/interfaces';
 import type { SelectedFile } from '../interfaces';
-import { useBusterChatContextSelector } from '@/context/Chats';
-import { useEffect } from 'react';
-import { useUnmount } from 'ahooks';
+import { useChatIndividual } from '@/context/Chats';
 import { useFileFallback } from './useFileFallback';
 
 export const useSubscribeIndividualChat = ({
@@ -12,23 +10,11 @@ export const useSubscribeIndividualChat = ({
   chatId?: string;
   defaultSelectedFile?: SelectedFile;
 }): IBusterChat | undefined => {
-  const chat: IBusterChat | undefined = useBusterChatContextSelector((x) => x.chats[chatId || '']);
-  const subscribeToChat = useBusterChatContextSelector((x) => x.subscribeToChat);
-  const unsubscribeFromChat = useBusterChatContextSelector((x) => x.unsubscribeFromChat);
+  const { chat } = useChatIndividual(chatId || '');
 
-  const { memoizedFallbackToChat } = useFileFallback({
-    defaultSelectedFile
-  });
+  const { memoizedFallbackToChat } = useFileFallback({ defaultSelectedFile });
 
   const selectedChat: IBusterChat | undefined = chatId ? chat : memoizedFallbackToChat;
 
-  useEffect(() => {
-    if (chatId) subscribeToChat({ chatId });
-  }, [chatId]);
-
-  useUnmount(() => {
-    if (chatId) unsubscribeFromChat({ chatId });
-  });
-
   return selectedChat;
 };

web/src/context/Chats/ChatProvider/ChatProvider.tsx

@@ -1,45 +1,22 @@
-import React, { useRef, useTransition } from 'react';
+import React from 'react';
 import {
   createContext,
   ContextSelector,
   useContextSelector
 } from '@fluentui/react-context-selector';
 import type { BusterChat } from '@/api/asset_interfaces';
-import type { IBusterChat, IBusterChatMessage } from '../interfaces';
-import { useChatSubscriptions } from './useChatSubscriptions';
 import { useChatAssosciations } from './useChatAssosciations';
 import { useChatSelectors } from './useChatSelectors';
 import { useChatUpdate } from './useChatUpdate';
 
 export const useBusterChat = () => {
-  const [isPending, startTransition] = useTransition();
-  const chatsRef = useRef<Record<string, IBusterChat>>({});
-  const chatsMessagesRef = useRef<Record<string, IBusterChatMessage>>({});
-
-  const chatSubscriptions = useChatSubscriptions({
-    chatsRef,
-    chatsMessagesRef,
-    startTransition
-  });
-
   const chatAssociations = useChatAssosciations();
 
-  const chatSelectors = useChatSelectors({
-    chatsRef,
-    chatsMessagesRef,
-    isPending
-  });
+  const chatSelectors = useChatSelectors();
 
-  const chatUpdate = useChatUpdate({
-    chatsRef,
-    chatsMessagesRef,
-    startTransition
-  });
+  const chatUpdate = useChatUpdate();
 
   return {
-    chats: chatsRef.current,
-    chatsMessages: chatsMessagesRef.current,
-    ...chatSubscriptions,
     ...chatAssociations,
     ...chatSelectors,
     ...chatUpdate

web/src/context/Chats/ChatProvider/index.ts

@@ -1 +1,6 @@
+import { useChatIndividual } from './useChatIndividual';
+import { useMessageIndividual } from './useMessageIndividual';
+
 export * from './ChatProvider';
+
+export { useChatIndividual, useMessageIndividual };

web/src/context/Chats/ChatProvider/useChatIndividual.ts

@@ -0,0 +1,31 @@
+import { updateChatToIChat } from '@/utils/chat';
+import { useSocketQueryEmitOn } from '@/api/buster_socket_query';
+import { queryKeys } from '@/api/query_keys';
+import { useQueryClient } from '@tanstack/react-query';
+
+export const useChatIndividual = (chatId: string) => {
+  const queryClient = useQueryClient();
+
+  const { data: chat, isFetched: isFetchedChat } = useSocketQueryEmitOn(
+    {
+      route: '/chats/get',
+      payload: { id: chatId }
+    },
+    '/chats/get:getChat',
+    queryKeys['chatsGetChat'](chatId),
+    (_currentData, newData) => {
+      const { iChat, iChatMessages } = updateChatToIChat(newData, false);
+      for (const message of iChatMessages) {
+        const options = queryKeys['chatsMessages'](message.id);
+        const queryKey = options.queryKey;
+        queryClient.setQueryData(queryKey, message);
+      }
+      return iChat;
+    }
+  );
+
+  return {
+    chat,
+    isFetchedChat
+  };
+};

web/src/context/Chats/ChatProvider/useChatSelectors.ts

@@ -1,47 +1,31 @@
 import { type MutableRefObject, useCallback } from 'react';
 import type { IBusterChat, IBusterChatMessage } from '../interfaces';
 import { useMemoizedFn } from 'ahooks';
+import { useQueryClient } from '@tanstack/react-query';
+import { queryKeys } from '@/api/query_keys';
+
+export const useChatSelectors = () => {
+  const queryClient = useQueryClient();
 
-export const useChatSelectors = ({
-  isPending,
-  chatsRef,
-  chatsMessagesRef
-}: {
-  isPending: boolean;
-  chatsRef: MutableRefObject<Record<string, IBusterChat>>;
-  chatsMessagesRef: MutableRefObject<Record<string, IBusterChatMessage>>;
-}) => {
   const getChatMemoized = useMemoizedFn((chatId: string): IBusterChat | undefined => {
-    return chatsRef.current[chatId];
+    const options = queryKeys['chatsGetChat'](chatId);
+    const queryKey = options.queryKey;
+    return queryClient.getQueryData<IBusterChat>(queryKey);
   });
 
-  const getChatMessages = useCallback(
-    (chatId: string): IBusterChatMessage[] => {
-      return getChatMessagesMemoized(chatId);
-    },
-    [chatsMessagesRef, isPending, chatsRef]
-  );
-
-  const getChatMessage = useCallback(
-    (messageId: string): IBusterChatMessage | undefined => {
-      return getChatMessageMemoized(messageId);
-    },
-    [chatsMessagesRef, isPending]
-  );
-
   const getChatMessagesMemoized = useMemoizedFn((chatId: string) => {
-    const chatMessageIds = chatsRef.current[chatId]?.messages || [];
-    return chatMessageIds.map((messageId) => chatsMessagesRef.current[messageId]);
+    const chatMessageIds = getChatMemoized(chatId)?.messages || [];
+    return chatMessageIds.map((messageId) => getChatMessageMemoized(messageId));
   });
 
   const getChatMessageMemoized = useMemoizedFn((messageId: string) => {
-    return chatsMessagesRef.current[messageId];
+    const options = queryKeys['chatsMessages'](messageId);
+    const queryKey = options.queryKey;
+    return queryClient.getQueryData<IBusterChatMessage>(queryKey);
   });
 
   return {
     getChatMemoized,
-    getChatMessages,
-    getChatMessage,
     getChatMessagesMemoized,
     getChatMessageMemoized
   };

web/src/context/Chats/ChatProvider/useChatSubscriptions.ts

@@ -1,59 +0,0 @@
-import { MutableRefObject } from 'react';
-import { useBusterWebSocket } from '../../BusterWebSocket';
-import { useMemoizedFn } from 'ahooks';
-import type { BusterChat } from '@/api/asset_interfaces';
-import type { IBusterChat, IBusterChatMessage } from '../interfaces';
-import { updateChatToIChat } from '@/utils/chat';
-
-export const useChatSubscriptions = ({
-  chatsRef,
-  chatsMessagesRef,
-  startTransition
-}: {
-  chatsRef: MutableRefObject<Record<string, IBusterChat>>;
-  chatsMessagesRef: MutableRefObject<Record<string, IBusterChatMessage>>;
-  startTransition: (fn: () => void) => void;
-}) => {
-  const busterSocket = useBusterWebSocket();
-
-  const _onGetChat = useMemoizedFn((chat: BusterChat): IBusterChat => {
-    const { iChat, iChatMessages } = updateChatToIChat(chat, false);
-
-    chatsRef.current[chat.id] = iChat;
-    chatsMessagesRef.current = {
-      ...chatsMessagesRef.current,
-      ...iChatMessages
-    };
-    startTransition(() => {
-      //just used to trigger UI update
-    });
-    return iChat;
-  });
-
-  const unsubscribeFromChat = useMemoizedFn(({ chatId }: { chatId: string }) => {
-    return busterSocket.emit({
-      route: '/chats/unsubscribe',
-      payload: {
-        id: chatId
-      }
-    });
-  });
-
-  const subscribeToChat = useMemoizedFn(({ chatId }: { chatId: string }) => {
-    return busterSocket.emitAndOnce({
-      emitEvent: {
-        route: '/chats/get',
-        payload: { id: chatId }
-      },
-      responseEvent: {
-        route: '/chats/get:getChat',
-        callback: _onGetChat
-      }
-    });
-  });
-
-  return {
-    unsubscribeFromChat,
-    subscribeToChat
-  };
-};

web/src/context/Chats/ChatProvider/useChatUpdate.ts

@@ -1,37 +1,39 @@
 import { useBusterWebSocket } from '@/context/BusterWebSocket';
 import { useMemoizedFn } from 'ahooks';
-import { type MutableRefObject } from 'react';
+import { useTransition, type MutableRefObject } from 'react';
 import type { IBusterChat, IBusterChatMessage } from '../interfaces';
+import { useSocketQueryMutation } from '@/api/buster_socket_query';
+import { useQueryClient } from '@tanstack/react-query';
+import { queryKeys } from '@/api/query_keys';
+import { create } from 'mutative';
 
-export const useChatUpdate = ({
-  chatsRef,
-  chatsMessagesRef,
-  startTransition
-}: {
-  chatsRef: MutableRefObject<Record<string, IBusterChat>>;
-  chatsMessagesRef: MutableRefObject<Record<string, IBusterChatMessage>>;
-  startTransition: (fn: () => void) => void;
-}) => {
-  const busterSocket = useBusterWebSocket();
+export const useChatUpdate = () => {
+  const [isPending, startTransition] = useTransition();
+  const queryClient = useQueryClient();
+
+  const { mutate: updateChat } = useSocketQueryMutation(
+    '/chats/update',
+    '/chats/update:updateChat',
+    null
+  );
 
   const onUpdateChat = useMemoizedFn(
     async (newChatConfig: Partial<IBusterChat> & { id: string }, saveToServer: boolean = false) => {
-      chatsRef.current[newChatConfig.id] = {
-        ...chatsRef.current[newChatConfig.id],
+      const options = queryKeys['chatsGetChat'](newChatConfig.id);
+      const queryKey = options.queryKey;
+      const currentData = queryClient.getQueryData<IBusterChat>(queryKey);
+      const iChat: IBusterChat = {
+        ...currentData!,
         ...newChatConfig
       };
+      queryClient.setQueryData(queryKey, iChat);
       startTransition(() => {
         //just used to trigger UI update
-
         if (saveToServer) {
-          const { title, is_favorited, id } = chatsRef.current[newChatConfig.id];
-          busterSocket.emit({
-            route: '/chats/update',
-            payload: {
-              id,
-              title,
-              is_favorited
-            }
+          updateChat({
+            id: iChat.id,
+            title: iChat.title,
+            is_favorited: iChat.is_favorited
           });
         }
       });
@@ -40,22 +42,17 @@ export const useChatUpdate = ({
 
   const onUpdateChatMessage = useMemoizedFn(
     async (newMessageConfig: Partial<IBusterChatMessage> & { id: string }) => {
-      chatsMessagesRef.current[newMessageConfig.id] = {
-        ...chatsMessagesRef.current[newMessageConfig.id],
+      const options = queryKeys['chatsMessages'](newMessageConfig.id);
+      const queryKey = options.queryKey;
+      const currentData = queryClient.getQueryData<IBusterChatMessage>(queryKey);
+      const iChatMessage: IBusterChatMessage = {
+        ...currentData!,
         ...newMessageConfig
       };
-      startTransition(() => {
-        //just used to trigger UI update
-      });
-    }
-  );
+      queryClient.setQueryData(queryKey, iChatMessage);
+
+      //TODO: update the message in the server
 
-  const onBulkSetChatMessages = useMemoizedFn(
-    (newMessagesConfig: Record<string, IBusterChatMessage>) => {
-      chatsMessagesRef.current = {
-        ...chatsMessagesRef.current,
-        ...newMessagesConfig
-      };
       startTransition(() => {
         //just used to trigger UI update
       });
@@ -64,7 +61,6 @@ export const useChatUpdate = ({
 
   return {
     onUpdateChat,
-    onUpdateChatMessage,
-    onBulkSetChatMessages
+    onUpdateChatMessage
   };
 };

web/src/context/Chats/ChatProvider/useMessageIndividual.ts

@@ -0,0 +1,8 @@
+import { queryKeys } from '@/api/query_keys';
+import { useQuery } from '@tanstack/react-query';
+
+export const useMessageIndividual = (messageId: string) => {
+  const options = queryKeys['chatsMessages'](messageId);
+  const { data: message } = useQuery({ ...options, enabled: false });
+  return message;
+};

web/src/context/Chats/NewChatProvider/NewChatProvider.tsx

@@ -7,19 +7,17 @@ import {
 import { useMemoizedFn } from 'ahooks';
 import type { BusterSearchResult, FileType } from '@/api/asset_interfaces';
 import { useBusterWebSocket } from '@/context/BusterWebSocket';
-import { useChatUpdateMessage } from './useChatUpdateMessage';
+import { useChatStreamMessage } from './useChatStreamMessage';
 
 export const useBusterNewChat = () => {
   const busterSocket = useBusterWebSocket();
 
   const {
     completeChatCallback,
-    startListeningForChatProgress,
-    stopListeningForChatProgress,
     stopChatCallback,
     initializeNewChatCallback,
     replaceMessageCallback
-  } = useChatUpdateMessage();
+  } = useChatStreamMessage();
 
   const onSelectSearchAsset = useMemoizedFn(async (asset: BusterSearchResult) => {
     console.log('select search asset');
@@ -38,8 +36,6 @@ export const useBusterNewChat = () => {
       metricId?: string;
       dashboardId?: string;
     }) => {
-      startListeningForChatProgress();
-
       await busterSocket.emitAndOnce({
         emitEvent: {
           route: '/chats/post',
@@ -56,14 +52,10 @@ export const useBusterNewChat = () => {
         }
       });
 
-      busterSocket
-        .once({
-          route: '/chats/post:complete',
-          callback: completeChatCallback
-        })
-        .then(() => {
-          stopListeningForChatProgress();
-        });
+      busterSocket.once({
+        route: '/chats/post:complete',
+        callback: completeChatCallback
+      });
     }
   );
 
@@ -95,7 +87,6 @@ export const useBusterNewChat = () => {
       messageId: string;
       chatId: string;
     }) => {
-      startListeningForChatProgress();
       replaceMessageCallback({
         prompt,
         messageId
@@ -114,14 +105,11 @@ export const useBusterNewChat = () => {
           callback: completeChatCallback
         }
       });
-      stopListeningForChatProgress();
     }
   );
 
   const onFollowUpChat = useMemoizedFn(
     async ({ prompt, chatId }: { prompt: string; chatId: string }) => {
-      startListeningForChatProgress();
-
       await busterSocket.emitAndOnce({
         emitEvent: {
           route: '/chats/post',
@@ -135,7 +123,6 @@ export const useBusterNewChat = () => {
           callback: completeChatCallback
         }
       });
-      stopListeningForChatProgress();
     }
   );
 
@@ -148,7 +135,6 @@ export const useBusterNewChat = () => {
           message_id: messageId
         }
       });
-      stopListeningForChatProgress();
       stopChatCallback(chatId);
     }
   );

web/src/context/Chats/NewChatProvider/useChatStreamMessage.ts

@@ -0,0 +1,202 @@
+import { useMemoizedFn } from 'ahooks';
+import { useBusterChatContextSelector } from '../ChatProvider';
+import { BusterChat, BusterChatMessage_text } from '@/api/asset_interfaces';
+import {
+  ChatEvent_GeneratingReasoningMessage,
+  ChatEvent_GeneratingResponseMessage,
+  ChatEvent_GeneratingTitle
+} from '@/api/buster_socket/chats';
+import { updateChatToIChat } from '@/utils/chat';
+import { useAutoAppendThought } from './useAutoAppendThought';
+import { useAppLayoutContextSelector } from '@/context/BusterAppLayout';
+import { BusterRoutes } from '@/routes';
+import { useSocketQueryOn } from '@/api/buster_socket_query';
+import { useRef } from 'react';
+import { IBusterChat, IBusterChatMessage } from '../interfaces';
+
+type ChatMessageResponseMessage = IBusterChatMessage['response_messages'][number] & {
+  index: number;
+};
+type ChatMessageReasoning = IBusterChatMessage['reasoning'][number] & {
+  index: number;
+};
+
+type ChatMessageResponseMessagesRef = Record<
+  //message_id
+  string,
+  //response_message_id
+  Record<string, ChatMessageResponseMessage>
+>;
+
+type ChatMessageReasoningMessageRef = Record<
+  //message_id
+  string,
+  //reasoning_message_id
+  Record<string, ChatMessageReasoning>
+>;
+
+export const useChatStreamMessage = () => {
+  // const busterSocket = useBusterWebSocket();
+  const getChatMessage = useBusterChatContextSelector((x) => x.getChatMessageMemoized);
+  const onChangePage = useAppLayoutContextSelector((x) => x.onChangePage);
+  const onUpdateChat = useBusterChatContextSelector((x) => x.onUpdateChat);
+  const onUpdateChatMessage = useBusterChatContextSelector((x) => x.onUpdateChatMessage);
+  const onToggleChatsModal = useAppLayoutContextSelector((s) => s.onToggleChatsModal);
+  const chatMessageResponseMessagesRef = useRef<ChatMessageResponseMessagesRef>({});
+  const chatMessageReasoningMessageRef = useRef<ChatMessageReasoningMessageRef>({});
+  const chatMessagesRef = useRef<Record<string, Partial<IBusterChatMessage>>>({});
+  const chatRef = useRef<Record<string, Partial<IBusterChat>>>({});
+
+  const { autoAppendThought } = useAutoAppendThought();
+
+  const completeChatCallback = useMemoizedFn((d: BusterChat) => {
+    const { iChat, iChatMessages } = updateChatToIChat(d, false);
+    //  onBulkSetChatMessages(iChatMessages);
+    onUpdateChat(iChat);
+  });
+
+  const stopChatCallback = useMemoizedFn((chatId: string) => {
+    onUpdateChatMessage({
+      id: chatId,
+      isCompletedStream: true
+    });
+  });
+
+  const initializeNewChatCallback = useMemoizedFn((d: BusterChat) => {
+    const { iChat, iChatMessages } = updateChatToIChat(d, true);
+    //onBulkSetChatMessages(iChatMessages);
+    onUpdateChat(iChat);
+    onChangePage({
+      route: BusterRoutes.APP_CHAT_ID,
+      chatId: iChat.id
+    });
+    onToggleChatsModal(false);
+  });
+
+  const replaceMessageCallback = useMemoizedFn(
+    ({ prompt, messageId }: { prompt: string; messageId: string }) => {
+      const currentMessage = getChatMessage(messageId);
+      const currentRequestMessage = currentMessage?.request_message!;
+
+      onUpdateChatMessage({
+        id: messageId,
+        request_message: {
+          ...currentRequestMessage,
+          request: prompt
+        },
+        reasoning: [],
+        response_messages: []
+      });
+    }
+  );
+
+  const _generatingTitleCallback = useMemoizedFn((_: null, newData: ChatEvent_GeneratingTitle) => {
+    const { chat_id, title, title_chunk, progress } = newData;
+    const isCompleted = progress === 'completed';
+    const currentTitle = chatRef.current[chat_id]?.title || '';
+    const newTitle = isCompleted ? title : currentTitle + title_chunk;
+    chatRef.current[chat_id] = {
+      ...chatRef.current[chat_id],
+      title: newTitle
+    };
+    onUpdateChat({
+      id: chat_id,
+      title: newTitle
+    });
+  });
+
+  useSocketQueryOn({
+    socketResponse: '/chats/post:generatingTitle',
+    callback: _generatingTitleCallback
+  });
+
+  const _generatingResponseMessageCallback = useMemoizedFn(
+    (_: null, d: ChatEvent_GeneratingResponseMessage) => {
+      const { message_id, response_message, chat_id } = d;
+      const responseMessageId = response_message.id;
+      const foundResponseMessage: undefined | ChatMessageResponseMessage =
+        chatMessageResponseMessagesRef.current[message_id][responseMessageId];
+      const isNewMessage = !foundResponseMessage;
+      const currentResponseMessages = chatMessagesRef.current[message_id]?.response_messages ?? [];
+
+      if (response_message.type === 'text') {
+        const existingMessage = (foundResponseMessage as BusterChatMessage_text)?.message || '';
+        const isStreaming = !!response_message.message_chunk;
+        if (isStreaming) {
+          response_message.message = existingMessage + response_message.message_chunk;
+        }
+      }
+
+      if (isNewMessage) {
+        chatMessageResponseMessagesRef.current[message_id][responseMessageId] = {
+          ...response_message,
+          index: currentResponseMessages.length
+        };
+      }
+
+      const messageToUse = chatMessageResponseMessagesRef.current[message_id][responseMessageId];
+
+      onUpdateChatMessage({
+        id: message_id,
+        response_messages: isNewMessage
+          ? [...currentResponseMessages, messageToUse]
+          : [
+              ...currentResponseMessages.slice(0, messageToUse.index),
+              messageToUse,
+              ...currentResponseMessages.slice(messageToUse.index + 1)
+            ]
+      });
+    }
+  );
+
+  useSocketQueryOn({
+    socketResponse: '/chats/post:generatingResponseMessage',
+    callback: _generatingResponseMessageCallback
+  });
+
+  const _generatingReasoningMessageCallback = useMemoizedFn(
+    (_: null, d: ChatEvent_GeneratingReasoningMessage) => {
+      const { message_id, reasoning, chat_id } = d;
+      const reasoningMessageId = reasoning.id;
+      const foundReasoningMessage: undefined | ChatMessageReasoning =
+        chatMessageReasoningMessageRef.current[message_id][reasoningMessageId];
+      const isNewMessage = !foundReasoningMessage;
+      const currentReasoning = chatMessagesRef.current[message_id]?.reasoning ?? [];
+
+      if (isNewMessage) {
+        chatMessageReasoningMessageRef.current[message_id][reasoningMessageId] = {
+          ...reasoning,
+          index: currentReasoning.length
+        };
+      }
+
+      const messageToUse = chatMessageReasoningMessageRef.current[message_id][reasoningMessageId];
+
+      const updatedReasoning = isNewMessage
+        ? [...currentReasoning, messageToUse]
+        : [
+            ...currentReasoning.slice(0, foundReasoningMessage.index),
+            reasoning,
+            ...currentReasoning.slice(foundReasoningMessage.index + 1)
+          ];
+
+      onUpdateChatMessage({
+        id: message_id,
+        reasoning: autoAppendThought(updatedReasoning, chat_id),
+        isCompletedStream: false
+      });
+    }
+  );
+
+  useSocketQueryOn({
+    socketResponse: '/chats/post:generatingReasoningMessage',
+    callback: _generatingReasoningMessageCallback
+  });
+
+  return {
+    initializeNewChatCallback,
+    completeChatCallback,
+    stopChatCallback,
+    replaceMessageCallback
+  };
+};

web/src/context/Chats/NewChatProvider/useChatUpdateMessage.ts

@@ -1,191 +0,0 @@
-import { useMemoizedFn } from 'ahooks';
-import { useBusterChatContextSelector } from '../ChatProvider';
-import { useBusterWebSocket } from '@/context/BusterWebSocket';
-import { BusterChat, BusterChatMessage_text } from '@/api/asset_interfaces';
-import {
-  ChatEvent_GeneratingReasoningMessage,
-  ChatEvent_GeneratingResponseMessage,
-  ChatEvent_GeneratingTitle
-} from '@/api/buster_socket/chats';
-import { updateChatToIChat } from '@/utils/chat';
-import { useAutoAppendThought } from './useAutoAppendThought';
-import { useAppLayoutContextSelector } from '@/context/BusterAppLayout';
-import { BusterRoutes } from '@/routes';
-
-export const useChatUpdateMessage = () => {
-  const busterSocket = useBusterWebSocket();
-  const onChangePage = useAppLayoutContextSelector((x) => x.onChangePage);
-  const onUpdateChat = useBusterChatContextSelector((x) => x.onUpdateChat);
-  const getChatMemoized = useBusterChatContextSelector((x) => x.getChatMemoized);
-  const onUpdateChatMessage = useBusterChatContextSelector((x) => x.onUpdateChatMessage);
-  const getChatMessageMemoized = useBusterChatContextSelector((x) => x.getChatMessageMemoized);
-  const onBulkSetChatMessages = useBusterChatContextSelector((x) => x.onBulkSetChatMessages);
-  const onToggleChatsModal = useAppLayoutContextSelector((s) => s.onToggleChatsModal);
-
-  const { autoAppendThought } = useAutoAppendThought();
-
-  const _generatingTitleCallback = useMemoizedFn((d: ChatEvent_GeneratingTitle) => {
-    const { chat_id, title, title_chunk } = d;
-    const isCompleted = d.progress === 'completed';
-    const currentChat = getChatMemoized(chat_id)!;
-    const currentTitle = currentChat?.title;
-    const newTitle = isCompleted ? title : currentTitle + title_chunk;
-    onUpdateChat({
-      ...currentChat,
-      title: newTitle
-    });
-  });
-
-  const _generatingResponseMessageCallback = useMemoizedFn(
-    (d: ChatEvent_GeneratingResponseMessage) => {
-      const { message_id, response_message } = d;
-      const currentResponseMessages = getChatMessageMemoized(message_id)?.response_messages ?? [];
-      const responseMessageId = response_message.id;
-      const foundResponseMessage = currentResponseMessages.find(
-        ({ id }) => id === responseMessageId
-      );
-      const isNewMessage = !foundResponseMessage;
-
-      if (response_message.type === 'text') {
-        const existingMessage = (foundResponseMessage as BusterChatMessage_text)?.message || '';
-        const isStreaming = !!response_message.message_chunk;
-        if (isStreaming) {
-          response_message.message = existingMessage + response_message.message_chunk;
-        }
-      }
-
-      onUpdateChatMessage({
-        id: message_id,
-        response_messages: isNewMessage
-          ? [...currentResponseMessages, response_message]
-          : currentResponseMessages.map((rm) =>
-              rm.id === responseMessageId ? response_message : rm
-            )
-      });
-    }
-  );
-
-  const _generatingReasoningMessageCallback = useMemoizedFn(
-    (d: ChatEvent_GeneratingReasoningMessage) => {
-      const { message_id, reasoning, chat_id } = d;
-      const currentReasoning = getChatMessageMemoized(message_id)?.reasoning ?? [];
-      const reasoningMessageId = reasoning.id;
-      const foundReasoningMessage = currentReasoning.find(({ id }) => id === reasoningMessageId);
-      const isNewMessage = !foundReasoningMessage;
-
-      const updatedReasoning = isNewMessage
-        ? [...currentReasoning, reasoning]
-        : currentReasoning.map((rm) => (rm.id === reasoningMessageId ? reasoning : rm));
-
-      onUpdateChatMessage({
-        id: message_id,
-        reasoning: autoAppendThought(updatedReasoning, chat_id),
-        isCompletedStream: false
-      });
-    }
-  );
-
-  const completeChatCallback = useMemoizedFn((d: BusterChat) => {
-    const { iChat, iChatMessages } = updateChatToIChat(d, false);
-    onBulkSetChatMessages(iChatMessages);
-    onUpdateChat(iChat);
-  });
-
-  const stopChatCallback = useMemoizedFn((chatId: string) => {
-    onUpdateChatMessage({
-      id: chatId,
-      isCompletedStream: true
-    });
-  });
-
-  const initializeNewChatCallback = useMemoizedFn((d: BusterChat) => {
-    const { iChat, iChatMessages } = updateChatToIChat(d, true);
-    onBulkSetChatMessages(iChatMessages);
-    onUpdateChat(iChat);
-    onChangePage({
-      route: BusterRoutes.APP_CHAT_ID,
-      chatId: iChat.id
-    });
-    onToggleChatsModal(false);
-  });
-
-  const replaceMessageCallback = useMemoizedFn(
-    ({ prompt, messageId }: { prompt: string; messageId: string }) => {
-      const currentMessage = getChatMessageMemoized(messageId);
-      const currentRequestMessage = currentMessage?.request_message!;
-
-      onUpdateChatMessage({
-        id: messageId,
-        request_message: {
-          ...currentRequestMessage,
-          request: prompt
-        },
-        reasoning: [],
-        response_messages: []
-      });
-    }
-  );
-
-  const listenForGeneratingTitle = useMemoizedFn(() => {
-    busterSocket.on({
-      route: '/chats/post:generatingTitle',
-      callback: _generatingTitleCallback
-    });
-  });
-
-  const stopListeningForGeneratingTitle = useMemoizedFn(() => {
-    busterSocket.off({
-      route: '/chats/post:generatingTitle',
-      callback: _generatingTitleCallback
-    });
-  });
-
-  const listenForGeneratingResponseMessage = useMemoizedFn(() => {
-    busterSocket.on({
-      route: '/chats/post:generatingResponseMessage',
-      callback: _generatingResponseMessageCallback
-    });
-  });
-
-  const stopListeningForGeneratingResponseMessage = useMemoizedFn(() => {
-    busterSocket.off({
-      route: '/chats/post:generatingResponseMessage',
-      callback: _generatingResponseMessageCallback
-    });
-  });
-
-  const listenForGeneratingReasoningMessage = useMemoizedFn(() => {
-    busterSocket.on({
-      route: '/chats/post:generatingReasoningMessage',
-      callback: _generatingReasoningMessageCallback
-    });
-  });
-
-  const stopListeningForGeneratingReasoningMessage = useMemoizedFn(() => {
-    busterSocket.off({
-      route: '/chats/post:generatingReasoningMessage',
-      callback: _generatingReasoningMessageCallback
-    });
-  });
-
-  const startListeningForChatProgress = useMemoizedFn(() => {
-    listenForGeneratingTitle();
-    listenForGeneratingResponseMessage();
-    listenForGeneratingReasoningMessage();
-  });
-
-  const stopListeningForChatProgress = useMemoizedFn(() => {
-    stopListeningForGeneratingTitle();
-    stopListeningForGeneratingResponseMessage();
-    stopListeningForGeneratingReasoningMessage();
-  });
-
-  return {
-    initializeNewChatCallback,
-    completeChatCallback,
-    startListeningForChatProgress,
-    stopListeningForChatProgress,
-    stopChatCallback,
-    replaceMessageCallback
-  };
-};

web/src/context/Chats/index.ts

@@ -1,8 +1,18 @@
 import { useBusterNewChatContextSelector } from './NewChatProvider';
-import { useBusterChatContextSelector } from './ChatProvider';
+import {
+  useBusterChatContextSelector,
+  useChatIndividual,
+  useMessageIndividual
+} from './ChatProvider';
 import { useBusterChatListByFilter } from './ChatListProvider';
 
 export * from './BusterChatProvider';
 
-export { useBusterNewChatContextSelector, useBusterChatContextSelector, useBusterChatListByFilter };
+export {
+  useBusterNewChatContextSelector,
+  useBusterChatContextSelector,
+  useBusterChatListByFilter,
+  useChatIndividual,
+  useMessageIndividual
+};
 export * from './interfaces';

web/src/utils/chat.ts

@@ -12,28 +12,23 @@ const chatUpgrader = (chat: BusterChat, { isNewChat }: { isNewChat: boolean }):
 const chatMessageUpgrader = (
   message: BusterChatMessage[],
   options?: { isCompletedStream: boolean; messageId: string }
-): Record<string, IBusterChatMessage> => {
+): IBusterChatMessage[] => {
   const lastMessageId = message[message.length - 1].id;
   const { isCompletedStream = true, messageId } = options || {};
   const optionMessageId = messageId || lastMessageId;
 
-  return message.reduce(
-    (acc, message) => {
-      if (message.id === optionMessageId) {
-        acc[message.id] = {
-          ...message,
-          isCompletedStream
-        };
-      } else {
-        acc[message.id] = {
-          ...message,
-          isCompletedStream: true
-        };
-      }
-      return acc;
-    },
-    {} as Record<string, IBusterChatMessage>
-  );
+  return message.map((message) => {
+    if (message.id === optionMessageId) {
+      return {
+        ...message,
+        isCompletedStream
+      };
+    }
+    return {
+      ...message,
+      isCompletedStream: true
+    };
+  });
 };
 
 export const updateChatToIChat = (chat: BusterChat, isNewChat: boolean) => {