better handling for black box messages

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

@@ -522,7 +522,8 @@ describe('updateReasoningMessage', () => {
   });
 
   it('should handle multiple updates and append text to existing reasoning message', () => {
-    const reasoning: BusterChatMessageReasoning_text = {
+    const messageId = 'test-message-id';
+    let reasoning: BusterChatMessageReasoning_text = {
       id: 'reasoning-1',
       type: 'text',
       message: '',
@@ -533,23 +534,46 @@ describe('updateReasoningMessage', () => {
     };
 
     // First update with "Hello"
-    let result = updateReasoningMessage('test-message-id', undefined, reasoning);
+    let result = updateReasoningMessage(messageId, undefined, reasoning);
     expect(
       (result.reasoning_messages['reasoning-1'] as BusterChatMessageReasoning_text).message
     ).toBe('Hello');
 
     // Second update with ", how"
-    reasoning.message_chunk = ', how';
-    result = updateReasoningMessage('test-message-id', result, reasoning);
+    reasoning = {
+      ...reasoning,
+      message_chunk: ', how'
+    };
+    result = updateReasoningMessage(messageId, result, reasoning);
     expect(
       (result.reasoning_messages['reasoning-1'] as BusterChatMessageReasoning_text).message
     ).toBe('Hello, how');
 
     // Third update with " are you doing today?"
-    reasoning.message_chunk = ' are you doing today?';
-    result = updateReasoningMessage('test-message-id', result, reasoning);
+    reasoning = {
+      ...reasoning,
+      message_chunk: ' are you doing today?'
+    };
+    result = updateReasoningMessage(messageId, result, reasoning);
     expect(
       (result.reasoning_messages['reasoning-1'] as BusterChatMessageReasoning_text).message
     ).toBe('Hello, how are you doing today?');
+
+    let reasoning2: BusterChatMessageReasoning_text = {
+      ...reasoning,
+      message_chunk: 'new reasoning message baby',
+      id: 'reasoning-2'
+    };
+
+    result = updateReasoningMessage(messageId, result, reasoning2);
+    expect(
+      (result.reasoning_messages['reasoning-1'] as BusterChatMessageReasoning_text).message
+    ).toBe('Hello, how are you doing today?');
+
+    expect(result.reasoning_message_ids).toContain('reasoning-2');
+    expect(result.reasoning_message_ids.length).toBe(2);
+    expect(
+      (result.reasoning_messages['reasoning-2'] as BusterChatMessageReasoning_text).message
+    ).toBe('new reasoning message baby');
   });
 });

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

@@ -16,55 +16,63 @@ export const useBlackBoxMessage = () => {
   const getChatMessageMemoized = useBusterChatContextSelector((x) => x.getChatMessageMemoized);
   const queryClient = useQueryClient();
 
-  const removeAutoThought = useMemoizedFn(({ messageId }: { messageId: string }) => {
-    console.log('removeAutoThought', messageId);
+  const clearTimeoutRef = useMemoizedFn((messageId: string) => {
     if (timeoutRef.current[messageId]) {
       clearTimeout(timeoutRef.current[messageId]);
       delete timeoutRef.current[messageId];
     }
+  });
+
+  const removeBlackBoxMessage = useMemoizedFn(({ messageId }: { messageId: string }) => {
+    console.log('removeBlackBoxMessage', messageId);
+    clearTimeoutRef(messageId);
 
     const options = queryKeys.chatsBlackBoxMessages(messageId);
     queryClient.setQueryData(options.queryKey, null);
   });
 
-  const addAutoThought = useMemoizedFn(({ messageId }: { messageId: string }) => {
+  const addBlackBoxMessage = useMemoizedFn(({ messageId }: { messageId: string }) => {
     const randomThought = getRandomThought();
     console.log(messageId, randomThought);
     const options = queryKeys.chatsBlackBoxMessages(messageId);
     queryClient.setQueryData(options.queryKey, randomThought);
   });
 
-  const checkAutoThought = useMemoizedFn(
+  const checkBlackBoxMessage = useMemoizedFn(
     (message: IBusterChatMessage, event: ChatEvent_GeneratingReasoningMessage) => {
       const isFinishedReasoningMessage = event.reasoning.status !== 'loading';
       if (isFinishedReasoningMessage) {
-        addAutoThought({ messageId: message.id });
+        clearTimeoutRef(message.id);
+        addBlackBoxMessage({ messageId: message.id });
         _loopAutoThought({ messageId: message.id });
       } else {
-        removeAutoThought({ messageId: message.id });
+        removeBlackBoxMessage({ messageId: message.id });
       }
     }
   );
 
   const _loopAutoThought = useMemoizedFn(async ({ messageId }: { messageId: string }) => {
-    const randomDelay = random(3000, 5000);
+    const randomDelay = random(5000, 5000);
     timeoutRef.current[messageId] = setTimeout(() => {
       const message = getChatMessageMemoized(messageId);
-      console.log('loopAutoThought', messageId, !!message);
       if (!message) return;
+      if (!timeoutRef.current[messageId]) return;
+
+      console.log('loopAutoThought', messageId, !!message);
+
       const isMessageCompletedStream = !!message?.isCompletedStream;
       const lastReasoningMessageId = last(message?.reasoning_message_ids) || '';
       const lastReasoningMessage = message?.reasoning_messages[lastReasoningMessageId];
       const isLastReasoningMessageCompleted = lastReasoningMessage?.status === 'completed';
 
       if (!isMessageCompletedStream && isLastReasoningMessageCompleted) {
-        addAutoThought({ messageId });
+        addBlackBoxMessage({ messageId });
         _loopAutoThought({ messageId });
       }
     }, randomDelay);
   });
 
-  return { checkAutoThought, removeAutoThought };
+  return { checkBlackBoxMessage, removeBlackBoxMessage };
 };
 
 const getRandomThought = (currentThought?: string): string => {

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

@@ -32,7 +32,7 @@ export const useChatStreamMessage = () => {
   const chatRefMessages = useRef<Record<string, IBusterChatMessage>>({});
   const [isPending, startTransition] = useTransition();
 
-  const { checkAutoThought, removeAutoThought } = useBlackBoxMessage();
+  const { checkBlackBoxMessage, removeBlackBoxMessage } = useBlackBoxMessage();
 
   const onUpdateChatMessageTransition = useMemoizedFn(
     (chatMessage: Parameters<typeof onUpdateChatMessage>[0]) => {
@@ -40,6 +40,7 @@ export const useChatStreamMessage = () => {
       const iChatMessage: IBusterChatMessage = create(currentChatMessage, (draft) => {
         Object.assign(draft || {}, chatMessage);
       })!;
+      chatRefMessages.current[chatMessage.id] = iChatMessage;
 
       onUpdateChatMessage(iChatMessage!);
 
@@ -65,7 +66,7 @@ export const useChatStreamMessage = () => {
     chatRef.current[iChat.id] = iChat;
     normalizeChatMessage(iChatMessages);
     onUpdateChat(iChat);
-    removeAutoThought({ messageId: iChat.message_ids[iChat.message_ids.length - 1] });
+    removeBlackBoxMessage({ messageId: iChat.message_ids[iChat.message_ids.length - 1] });
   });
 
   const stopChatCallback = useMemoizedFn((chatId: string) => {
@@ -127,12 +128,13 @@ export const useChatStreamMessage = () => {
   const _generatingReasoningMessageCallback = useMemoizedFn(
     (_: null, d: ChatEvent_GeneratingReasoningMessage) => {
       const { message_id, reasoning } = d;
-      const currentMessage = chatRefMessages.current[message_id];
-      const updatedMessage = updateReasoningMessage(message_id, currentMessage, reasoning);
+      const updatedMessage = updateReasoningMessage(
+        message_id,
+        chatRefMessages.current[message_id],
+        reasoning
+      );
 
-      checkAutoThought(updatedMessage, d);
-
-      //TRIGGER
+      checkBlackBoxMessage(updatedMessage, d); //we only trigger black box message if the reasoning message is completed
 
       onUpdateChatMessageTransition({
         id: message_id,

web/src/controllers/ReasoningController/ReasoningMessages/ReasoningBlackBoxMessage.tsx

@@ -23,7 +23,7 @@ export const BlackBoxMessage: React.FC<{ messageId: string }> = React.memo(({ me
     </BarContainer>;
   }
 
-  return null;
+  return <span>no black box message?</span>;
 });
 
 BlackBoxMessage.displayName = 'BlackBoxMessage';