move some functionality to a helper for testing

web/src/context/Chats/NewChatProvider/chatStreamMessageHelper.test.ts

@@ -0,0 +1,121 @@
+import { initializeOrUpdateMessage, updateChatTitle } from './chatStreamMessageHelper';
+import { IBusterChatMessage, IBusterChat } from '../interfaces';
+
+describe('initializeOrUpdateMessage', () => {
+  it('should initialize a new message when currentMessage is undefined', () => {
+    const messageId = 'test-id';
+    const updateFn = (draft: IBusterChatMessage) => {
+      if (draft.request_message) {
+        draft.request_message.request = 'test request';
+      }
+    };
+
+    const result = initializeOrUpdateMessage(messageId, undefined, updateFn);
+
+    expect(result.id).toBe(messageId);
+    expect(result.isCompletedStream).toBe(false);
+    expect(result.request_message?.request).toBe('test request');
+    expect(result.created_at).toBeDefined();
+    expect(result.final_reasoning_message).toBeNull();
+  });
+
+  it('should update an existing message', () => {
+    const messageId = 'test-id';
+    const currentMessage: IBusterChatMessage = {
+      id: messageId,
+      isCompletedStream: false,
+      request_message: {
+        request: 'original request',
+        sender_id: 'user1',
+        sender_name: 'Test User',
+        sender_avatar: null
+      },
+      response_message_ids: [],
+      reasoning_message_ids: [],
+      response_messages: {},
+      reasoning_messages: {},
+      created_at: new Date().toISOString(),
+      final_reasoning_message: null
+    };
+
+    const updateFn = (draft: IBusterChatMessage) => {
+      if (draft.request_message) {
+        draft.request_message.request = 'updated request';
+      }
+    };
+
+    const result = initializeOrUpdateMessage(messageId, currentMessage, updateFn);
+
+    expect(result.id).toBe(messageId);
+    expect(result.request_message?.request).toBe('updated request');
+  });
+});
+
+const mockChat: IBusterChat = {
+  id: 'test-chat-id',
+  title: 'Initial Title',
+  isNewChat: false,
+  is_favorited: false,
+  message_ids: [],
+  created_at: new Date().toISOString(),
+  updated_at: new Date().toISOString(),
+  created_by: 'test-user',
+  created_by_id: 'test-user-id',
+  created_by_name: 'Test User',
+  created_by_avatar: null
+};
+
+describe('updateChatTitle', () => {
+  it('should update title with a chunk when progress is not completed', () => {
+    const event = {
+      chat_id: 'test-chat-id',
+      title: 'Final Title',
+      title_chunk: ' New',
+      progress: 'in_progress' as const
+    };
+
+    const result = updateChatTitle(mockChat, event);
+    expect(result.title).toBe('Initial Title New');
+  });
+
+  it('should set the final title when progress is completed', () => {
+    const event = {
+      chat_id: 'test-chat-id',
+      title: 'Final Title',
+      title_chunk: '',
+      progress: 'completed' as const
+    };
+
+    const result = updateChatTitle(mockChat, event);
+    expect(result.title).toBe('Final Title');
+  });
+
+  it('should handle chat with empty initial title', () => {
+    const chatWithoutTitle: IBusterChat = {
+      ...mockChat,
+      title: ''
+    };
+
+    const event = {
+      chat_id: 'test-chat-id',
+      title: 'Final Title',
+      title_chunk: ' New',
+      progress: 'in_progress' as const
+    };
+
+    const result = updateChatTitle(chatWithoutTitle, event);
+    expect(result.title).toBe(' New');
+  });
+
+  it('should not modify title when title and title_chunk are empty', () => {
+    const event = {
+      chat_id: 'test-chat-id',
+      title: '',
+      title_chunk: '',
+      progress: 'completed' as const
+    };
+
+    const result = updateChatTitle(mockChat, event);
+    expect(result.title).toBe('Initial Title');
+  });
+});

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

@@ -0,0 +1,43 @@
+import { create } from 'mutative';
+import { IBusterChat, IBusterChatMessage } from '../interfaces';
+import { ChatEvent_GeneratingTitle } from '@/api/buster_socket/chats';
+
+const createInitialMessage = (messageId: string): IBusterChatMessage => ({
+  id: messageId,
+  isCompletedStream: false,
+  request_message: {
+    request: '',
+    sender_id: '',
+    sender_name: '',
+    sender_avatar: null
+  },
+  response_message_ids: [],
+  reasoning_message_ids: [],
+  response_messages: {},
+  reasoning_messages: {},
+  created_at: new Date().toISOString(),
+  final_reasoning_message: null
+});
+
+export const initializeOrUpdateMessage = (
+  messageId: string,
+  currentMessage: IBusterChatMessage | undefined,
+  updateFn: (draft: IBusterChatMessage) => void
+) => {
+  return create(currentMessage || createInitialMessage(messageId), (draft) => {
+    updateFn(draft);
+  });
+};
+
+export const updateChatTitle = (
+  currentChat: IBusterChat,
+  event: ChatEvent_GeneratingTitle
+): IBusterChat => {
+  const { chat_id, title, title_chunk, progress } = event;
+  const isCompleted = progress === 'completed';
+  const currentTitle = currentChat.title || '';
+  const newTitle = isCompleted ? title : currentTitle + title_chunk;
+  return create(currentChat, (draft) => {
+    if (newTitle) draft.title = newTitle;
+  });
+};

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

@@ -22,6 +22,7 @@ import { IBusterChat, IBusterChatMessage } from '../interfaces';
 import { queryKeys } from '@/api/query_keys';
 import { useQueryClient } from '@tanstack/react-query';
 import { create } from 'mutative';
+import { initializeOrUpdateMessage, updateChatTitle } from './chatStreamMessageHelper';
 
 export const useChatStreamMessage = () => {
   const queryClient = useQueryClient();
@@ -50,17 +51,6 @@ export const useChatStreamMessage = () => {
     }
   );
 
-  const initializeOrUpdateMessage = useMemoizedFn(
-    (messageId: string, updateFn: (draft: IBusterChatMessage) => void) => {
-      const currentMessage = chatRefMessages.current[messageId];
-      const updatedMessage = create(currentMessage || {}, (draft) => {
-        updateFn(draft);
-      });
-      chatRefMessages.current[messageId] = updatedMessage;
-      onUpdateChatMessage(updatedMessage);
-    }
-  );
-
   const normalizeChatMessage = useMemoizedFn(
     (iChatMessages: Record<string, IBusterChatMessage>) => {
       for (const message of Object.values(iChatMessages)) {
@@ -93,7 +83,6 @@ export const useChatStreamMessage = () => {
     chatRef.current = create(chatRef.current, (draft) => {
       draft[iChat.id] = iChat;
     });
-
     normalizeChatMessage(iChatMessages);
     onUpdateChat(iChat);
     onChangePage({
@@ -119,32 +108,27 @@ export const useChatStreamMessage = () => {
   );
 
   const _generatingTitleCallback = useMemoizedFn((_: null, newData: ChatEvent_GeneratingTitle) => {
-    const { chat_id, title, title_chunk, progress } = newData;
+    const { chat_id } = newData;
-    const isCompleted = progress === 'completed';
+    const updatedChat = updateChatTitle(chatRef.current[chat_id], newData);
-    const currentTitle = chatRef.current[chat_id]?.title || '';
+    chatRef.current[chat_id] = updatedChat;
-    const newTitle = isCompleted ? title : currentTitle + title_chunk;
+    onUpdateChat(updatedChat);
-    chatRef.current = create(chatRef.current, (draft) => {
-      if (newTitle) draft[chat_id].title = newTitle;
-    });
-    onUpdateChat({
-      id: chat_id,
-      title: newTitle
-    });
   });
 
   const _generatingResponseMessageCallback = useMemoizedFn(
     (_: null, d: ChatEvent_GeneratingResponseMessage) => {
-      const { message_id, response_message, chat_id } = d;
+      const { message_id, response_message } = d;
 
       if (!response_message?.id) return;
 
       const responseMessageId = response_message.id;
-      const existingMessage =
+      const existingResponseMessage =
         chatRefMessages.current[message_id]?.response_messages?.[responseMessageId];
-      const isNewMessage = !existingMessage;
+      const isNewResponseMessage = !existingResponseMessage;
 
-      if (isNewMessage) {
+      let currentMessage = chatRefMessages.current[message_id];
-        initializeOrUpdateMessage(message_id, (draft) => {
+
+      if (isNewResponseMessage) {
+        currentMessage = initializeOrUpdateMessage(message_id, currentMessage, (draft) => {
           if (!draft.response_messages) {
             draft.response_messages = {};
           }
@@ -157,11 +141,12 @@ export const useChatStreamMessage = () => {
       }
 
       if (response_message.type === 'text') {
-        const existingResponseMessageText = existingMessage as BusterChatResponseMessage_text;
+        const existingResponseMessageText =
+          existingResponseMessage as BusterChatResponseMessage_text;
         const isStreaming =
           response_message.message_chunk !== undefined && response_message.message_chunk !== null;
 
-        initializeOrUpdateMessage(message_id, (draft) => {
+        currentMessage = initializeOrUpdateMessage(message_id, currentMessage, (draft) => {
           const responseMessage = draft.response_messages?.[responseMessageId];
           if (!responseMessage) return;
           const messageText = responseMessage as BusterChatMessageReasoning_text;
@@ -176,7 +161,6 @@ export const useChatStreamMessage = () => {
         });
       }
 
-      const currentMessage = chatRefMessages.current[message_id];
       onUpdateChatMessageTransition({
         id: message_id,
         response_messages: currentMessage?.response_messages,
@@ -190,12 +174,13 @@ export const useChatStreamMessage = () => {
       const { message_id, reasoning, chat_id } = d;
 
       const reasoningMessageId = reasoning.id;
-      const existingMessage =
+      const existingReasoningMessage =
         chatRefMessages.current[message_id]?.reasoning_messages?.[reasoningMessageId];
-      const isNewMessage = !existingMessage;
+      const isNewReasoningMessage = !existingReasoningMessage;
+      let currentMessage = chatRefMessages.current[message_id];
 
-      if (isNewMessage) {
+      if (isNewReasoningMessage) {
-        initializeOrUpdateMessage(message_id, (draft) => {
+        currentMessage = initializeOrUpdateMessage(message_id, currentMessage, (draft) => {
           if (!draft.reasoning_messages) {
             draft.reasoning_messages = {};
           }
@@ -209,11 +194,12 @@ export const useChatStreamMessage = () => {
 
       switch (reasoning.type) {
         case 'text': {
-          const existingReasoningMessageText = existingMessage as BusterChatMessageReasoning_text;
+          const existingReasoningMessageText =
+            existingReasoningMessage as BusterChatMessageReasoning_text;
           const isStreaming =
             reasoning.message_chunk !== null || reasoning.message_chunk !== undefined;
 
-          initializeOrUpdateMessage(message_id, (draft) => {
+          currentMessage = initializeOrUpdateMessage(message_id, currentMessage, (draft) => {
             const reasoningMessage = draft.reasoning_messages?.[reasoningMessageId];
             if (!reasoningMessage) return;
             const messageText = reasoningMessage as BusterChatMessageReasoning_text;
@@ -230,9 +216,10 @@ export const useChatStreamMessage = () => {
           break;
         }
         case 'files': {
-          const existingReasoningMessageFiles = existingMessage as BusterChatMessageReasoning_files;
+          const existingReasoningMessageFiles =
+            existingReasoningMessage as BusterChatMessageReasoning_files;
 
-          initializeOrUpdateMessage(message_id, (draft) => {
+          currentMessage = initializeOrUpdateMessage(message_id, currentMessage, (draft) => {
             const reasoningMessage = draft.reasoning_messages?.[reasoningMessageId];
             if (!reasoningMessage) return;
 
@@ -286,7 +273,7 @@ export const useChatStreamMessage = () => {
           break;
         }
         case 'pills': {
-          initializeOrUpdateMessage(message_id, (draft) => {
+          currentMessage = initializeOrUpdateMessage(message_id, currentMessage, (draft) => {
             if (!draft.reasoning_messages?.[reasoningMessageId]) return;
             draft.reasoning_messages[reasoningMessageId] = reasoning;
           });
@@ -299,7 +286,6 @@ export const useChatStreamMessage = () => {
         }
       }
 
-      const currentMessage = chatRefMessages.current[message_id];
       onUpdateChatMessageTransition({
         id: message_id,
         reasoning_messages: currentMessage?.reasoning_messages,