upate replace message logic

web/src/api/asset_interfaces/chat/chatMessageInterfaces.ts

@@ -3,7 +3,7 @@ import type { FileType, ThoughtFileType } from './config';
 
 export type BusterChatMessage = {
   id: string;
-  request_message: BusterChatMessageRequest;
+  request_message: BusterChatMessageRequest | null;
   response_message_ids: string[];
   response_messages: Record<string, BusterChatMessageResponse>;
   reasoning_message_ids: string[];

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

@@ -153,7 +153,7 @@ export const useDeleteChat = () => {
   });
 };
 
-export const useGetChatMemoized = () => {
+export const useGetChatMessageMemoized = () => {
   const queryClient = useQueryClient();
 
   const getChatMessageMemoized = useMemoizedFn((messageId: string) => {
@@ -165,6 +165,18 @@ export const useGetChatMemoized = () => {
   return getChatMessageMemoized;
 };
 
+export const useGetChatMemoized = () => {
+  const queryClient = useQueryClient();
+
+  const getChatMemoized = useMemoizedFn((chatId: string) => {
+    const options = queryKeys.chatsGetChat(chatId);
+    const queryKey = options.queryKey;
+    return queryClient.getQueryData<IBusterChat>(queryKey);
+  });
+
+  return getChatMemoized;
+};
+
 export const useGetChatMessage = <TData = IBusterChatMessage>(
   messageId: string,
   selector?: (message: IBusterChatMessage) => TData

web/src/context/Chats/NewChatProvider.tsx

@@ -6,16 +6,18 @@ import { useMemoizedFn } from '@/hooks';
 import type { BusterSearchResult, FileType } from '@/api/asset_interfaces';
 import { useBusterWebSocket } from '@/context/BusterWebSocket';
 import { useChatStreamMessage } from './useChatStreamMessage';
+import { useGetChatMemoized, useGetChatMessageMemoized } from '@/api/buster_rest/chats';
+import { useChatUpdate } from './useChatUpdate';
+import { create } from 'mutative';
 
 export const useBusterNewChat = () => {
   const busterSocket = useBusterWebSocket();
+  const getChatMessageMemoized = useGetChatMessageMemoized();
+  const getChatMemoized = useGetChatMemoized();
+  const { onUpdateChat, onUpdateChatMessage } = useChatUpdate();
 
-  const {
-    completeChatCallback,
-    stopChatCallback,
-    initializeNewChatCallback,
-    replaceMessageCallback
-  } = useChatStreamMessage();
+  const { completeChatCallback, stopChatCallback, initializeNewChatCallback } =
+    useChatStreamMessage();
 
   const onSelectSearchAsset = useMemoizedFn(async (asset: BusterSearchResult) => {
     await new Promise((resolve) => setTimeout(resolve, 1000));
@@ -84,10 +86,31 @@ export const useBusterNewChat = () => {
       messageId: string;
       chatId: string;
     }) => {
-      replaceMessageCallback({
-        prompt,
-        messageId
+      const currentChat = getChatMemoized(chatId);
+      const currentMessage = getChatMessageMemoized(messageId);
+      const currentRequestMessage = currentMessage?.request_message!;
+      onUpdateChatMessage({
+        id: messageId,
+        request_message: create(currentRequestMessage, (draft) => {
+          draft.request = prompt;
+        }),
+        reasoning_message_ids: [],
+        response_message_ids: [],
+        isCompletedStream: false
       });
+
+      const messageIndex = currentChat?.message_ids.findIndex(
+        (messageId) => messageId === messageId
+      );
+
+      if (messageIndex && messageIndex !== -1) {
+        const updatedMessageIds = currentChat?.message_ids.slice(0, messageIndex + 1);
+        onUpdateChat({
+          id: chatId,
+          message_ids: updatedMessageIds
+        });
+      }
+
       await busterSocket.emitAndOnce({
         emitEvent: {
           route: '/chats/post',

web/src/context/Chats/useBlackBoxMessage.ts

@@ -9,11 +9,11 @@ import { IBusterChatMessage } from '@/api/asset_interfaces/chat';
 import { ChatEvent_GeneratingReasoningMessage } from '@/api/buster_socket/chats';
 import { useQueryClient } from '@tanstack/react-query';
 import { queryKeys } from '@/api/query_keys';
-import { useGetChatMemoized } from '@/api/buster_rest/chats';
+import { useGetChatMessageMemoized } from '@/api/buster_rest/chats';
 
 export const useBlackBoxMessage = () => {
   const timeoutRef = useRef<Record<string, ReturnType<typeof setTimeout>>>({});
-  const getChatMessageMemoized = useGetChatMemoized();
+  const getChatMessageMemoized = useGetChatMessageMemoized();
   const queryClient = useQueryClient();
 
   const clearTimeoutRef = useMemoizedFn((messageId: string) => {

web/src/context/Chats/useChatStreamMessage.ts

@@ -21,14 +21,12 @@ import {
   updateResponseMessage,
   updateReasoningMessage
 } from './chatStreamMessageHelper';
-import { useGetChatMemoized } from '@/api/buster_rest/chats';
 import { useChatUpdate } from './useChatUpdate';
-import { prefetchGetMetricDataClient, prefetchGetMetric } from '@/api/buster_rest/metrics';
+import { prefetchGetMetricDataClient } from '@/api/buster_rest/metrics';
 
 export const useChatStreamMessage = () => {
   const queryClient = useQueryClient();
   const onChangePage = useAppLayoutContextSelector((x) => x.onChangePage);
-  const getChatMessageMemoized = useGetChatMemoized();
   const { onUpdateChat, onUpdateChatMessage } = useChatUpdate();
   const chatRef = useRef<Record<string, IBusterChat>>({});
   const chatRefMessages = useRef<Record<string, IBusterChatMessage>>({});
@@ -110,21 +108,6 @@ export const useChatStreamMessage = () => {
     }
   });
 
-  const replaceMessageCallback = useMemoizedFn(
-    ({ prompt, messageId }: { prompt: string; messageId: string }) => {
-      const currentMessage = getChatMessageMemoized(messageId);
-      const currentRequestMessage = currentMessage?.request_message!;
-      onUpdateChatMessage({
-        id: messageId,
-        request_message: create(currentRequestMessage, (draft) => {
-          draft.request = prompt;
-        }),
-        reasoning_message_ids: [],
-        response_message_ids: []
-      });
-    }
-  );
-
   const _generatingTitleCallback = useMemoizedFn((_: null, newData: ChatEvent_GeneratingTitle) => {
     const { chat_id } = newData;
     const currentChat = chatRef.current[chat_id];
@@ -189,7 +172,6 @@ export const useChatStreamMessage = () => {
   return {
     initializeNewChatCallback,
     completeChatCallback,
-    stopChatCallback,
-    replaceMessageCallback
+    stopChatCallback
   };
 };

web/src/layouts/ChatLayout/ChatContainer/ChatContent/ChatMessageBlock.tsx

@@ -7,14 +7,22 @@ export const ChatMessageBlock: React.FC<{
   messageId: string;
   chatId: string;
 }> = React.memo(({ messageId, chatId }) => {
+  const messageExists = useGetChatMessage(messageId, (message) => message?.id);
   const requestMessage = useGetChatMessage(messageId, (message) => message?.request_message);
   const isCompletedStream = useGetChatMessage(messageId, (x) => x?.isCompletedStream);
 
-  if (!requestMessage) return null;
+  if (!messageExists) return null;
 
   return (
     <div className={'flex flex-col space-y-3.5 py-2 pr-3 pl-4'} id={messageId}>
-      <ChatUserMessage requestMessage={requestMessage} />
+      {requestMessage && (
+        <ChatUserMessage
+          isCompletedStream={isCompletedStream!}
+          chatId={chatId}
+          messageId={messageId}
+          requestMessage={requestMessage}
+        />
+      )}
       <ChatResponseMessages
         isCompletedStream={isCompletedStream!}
         messageId={messageId}

web/src/layouts/ChatLayout/ChatContainer/ChatContent/ChatUserMessage.tsx

@@ -1,20 +1,145 @@
+'use client';
+
 import type { BusterChatMessageRequest } from '@/api/asset_interfaces';
-import React from 'react';
+import React, { useState, useRef } from 'react';
 import { Paragraph } from '@/components/ui/typography';
 import { MessageContainer } from './MessageContainer';
+import { Tooltip } from '@/components/ui/tooltip';
+import { cn } from '@/lib/classMerge';
+import { PenWriting, Copy, Check } from '@/components/ui/icons';
+import { Button } from '@/components/ui/buttons';
+import { useBusterNotifications } from '@/context/BusterNotifications';
+import { useMemoizedFn } from '@/hooks';
+import { InputTextArea } from '@/components/ui/inputs/InputTextArea';
+import { useBusterNewChatContextSelector } from '@/context/Chats';
 
-export const ChatUserMessage: React.FC<{ requestMessage: BusterChatMessageRequest }> = React.memo(
-  ({ requestMessage }) => {
-    if (!requestMessage) return null;
+export const ChatUserMessage: React.FC<{
+  messageId: string;
+  chatId: string;
+  isCompletedStream: boolean;
+  requestMessage: NonNullable<BusterChatMessageRequest>;
+}> = React.memo(({ messageId, chatId, isCompletedStream, requestMessage }) => {
+  const [isTooltipOpen, setIsTooltipOpen] = useState(false);
+  const [isEditing, setIsEditing] = useState(false);
 
-    const { sender_avatar, sender_id, sender_name, request } = requestMessage;
+  const { sender_avatar, sender_id, sender_name, request } = requestMessage;
 
-    return (
-      <MessageContainer senderName={sender_name} senderId={sender_id} senderAvatar={sender_avatar}>
-        <Paragraph>{request}</Paragraph>
-      </MessageContainer>
-    );
-  }
-);
+  const onSetIsEditing = useMemoizedFn((isEditing: boolean) => {
+    setIsEditing(isEditing);
+    setIsTooltipOpen(false);
+  });
+
+  return (
+    <MessageContainer
+      senderName={sender_name}
+      senderId={sender_id}
+      senderAvatar={sender_avatar}
+      onMouseEnter={() => setIsTooltipOpen(true)}
+      onMouseLeave={() => setIsTooltipOpen(false)}>
+      {isEditing ? (
+        <EditMessage
+          messageId={messageId}
+          chatId={chatId}
+          requestMessage={requestMessage}
+          onSetIsEditing={onSetIsEditing}
+        />
+      ) : (
+        <>
+          <Paragraph>{request}</Paragraph>
+          {isCompletedStream && (
+            <RequestMessageTooltip
+              isTooltipOpen={isTooltipOpen}
+              requestMessage={requestMessage}
+              setIsEditing={setIsEditing}
+            />
+          )}
+        </>
+      )}
+    </MessageContainer>
+  );
+});
 
 ChatUserMessage.displayName = 'ChatUserMessage';
+
+const RequestMessageTooltip: React.FC<{
+  isTooltipOpen: boolean;
+  requestMessage: NonNullable<BusterChatMessageRequest>;
+  setIsEditing: (isEditing: boolean) => void;
+}> = React.memo(({ isTooltipOpen, requestMessage, setIsEditing }) => {
+  const { openSuccessMessage } = useBusterNotifications();
+
+  const onCopy = useMemoizedFn(() => {
+    navigator.clipboard.writeText(requestMessage.request);
+    openSuccessMessage('Copied to clipboard');
+  });
+
+  const onEdit = useMemoizedFn(() => {
+    setIsEditing(true);
+  });
+
+  return (
+    <div
+      className={cn(
+        'absolute right-1 -bottom-0 translate-y-full transform',
+        'bg-background z-50 rounded border shadow',
+        'transition-all duration-200',
+        isTooltipOpen ? 'scale-100 opacity-100' : 'scale-95 opacity-0'
+      )}>
+      <Tooltip title={'Edit'} side={'bottom'}>
+        <Button
+          prefix={<PenWriting />}
+          className="hover:bg-item-select!"
+          variant={'ghost'}
+          onClick={onEdit}
+        />
+      </Tooltip>
+      <Tooltip title={'Copy'} side={'bottom'}>
+        <Button
+          prefix={<Copy />}
+          className="hover:bg-item-select!"
+          variant={'ghost'}
+          onClick={onCopy}
+        />
+      </Tooltip>
+    </div>
+  );
+});
+
+RequestMessageTooltip.displayName = 'RequestMessageTooltip';
+
+const EditMessage: React.FC<{
+  requestMessage: NonNullable<BusterChatMessageRequest>;
+  onSetIsEditing: (isEditing: boolean) => void;
+  messageId: string;
+  chatId: string;
+}> = React.memo(({ requestMessage, onSetIsEditing, messageId, chatId }) => {
+  const [prompt, setPrompt] = useState(requestMessage.request);
+  const onReplaceMessageInChat = useBusterNewChatContextSelector((x) => x.onReplaceMessageInChat);
+
+  const onSave = useMemoizedFn((text: string) => {
+    onReplaceMessageInChat({
+      chatId,
+      messageId,
+      prompt
+    });
+    onSetIsEditing(false);
+  });
+
+  return (
+    <div className="-mt-1 flex flex-col space-y-2">
+      <InputTextArea
+        autoResize={{ minRows: 3, maxRows: 10 }}
+        value={prompt}
+        onChange={(e) => setPrompt(e.target.value)}
+      />
+      <div className="flex justify-end space-x-2">
+        <Button variant={'ghost'} onClick={() => onSetIsEditing(false)}>
+          Cancel
+        </Button>
+        <Button variant={'black'} onClick={() => onSave(prompt)}>
+          Submit
+        </Button>
+      </div>
+    </div>
+  );
+});

web/src/layouts/ChatLayout/ChatContainer/ChatContent/MessageContainer.tsx

@@ -1,24 +1,37 @@
 import { Avatar } from '@/components/ui/avatar';
 import { cn } from '@/lib/classMerge';
-import React from 'react';
+import React, { forwardRef } from 'react';
 
-export const MessageContainer: React.FC<{
+interface MessageContainerProps {
   children: React.ReactNode;
   senderName?: string;
   senderId?: string;
   senderAvatar?: string | null;
   className?: string;
-}> = React.memo(({ children, senderName, senderId, senderAvatar, className = '' }) => {
-  return (
-    <div className={'flex w-full space-x-2 overflow-hidden'}>
-      {senderName ? (
-        <Avatar size={24} name={senderName} image={senderAvatar || ''} useToolTip={true} />
-      ) : (
-        <Avatar size={24} />
-      )}
-      <div className={cn('mt-1 px-1', className)}>{children}</div>
-    </div>
-  );
-});
+  onMouseEnter?: () => void;
+  onMouseLeave?: () => void;
+}
+
+export const MessageContainer = forwardRef<HTMLDivElement, MessageContainerProps>(
+  (
+    { children, senderName, senderId, senderAvatar, className = '', onMouseEnter, onMouseLeave },
+    ref
+  ) => {
+    return (
+      <div
+        ref={ref}
+        className={'flex w-full space-x-2'}
+        onMouseEnter={onMouseEnter}
+        onMouseLeave={onMouseLeave}>
+        {senderName ? (
+          <Avatar size={24} name={senderName} image={senderAvatar || ''} useToolTip={true} />
+        ) : (
+          <Avatar size={24} />
+        )}
+        <div className={cn('relative mt-1 w-full px-1', className)}>{children}</div>
+      </div>
+    );
+  }
+);
 
 MessageContainer.displayName = 'MessageContainer';

web/src/layouts/ChatLayout/ChatContext/useAutoChangeLayout.ts

@@ -1,12 +1,10 @@
 'use client';
 
-import { useGetChatMemoized, useGetChatMessage } from '@/api/buster_rest/chats';
+import { useGetChatMessageMemoized, useGetChatMessage } from '@/api/buster_rest/chats';
 import type { SelectedFile } from '../interfaces';
 import { useEffect, useRef } from 'react';
 import findLast from 'lodash/findLast';
 import { BusterChatResponseMessage_file } from '@/api/asset_interfaces/chat';
-import { useAppLayoutContextSelector } from '@/context/BusterAppLayout';
-import { BusterRoutes } from '@/routes';
 
 export const useAutoChangeLayout = ({
   lastMessageId,
@@ -24,7 +22,7 @@ export const useAutoChangeLayout = ({
     lastMessageId,
     (x) => x?.reasoning_message_ids?.length || 0
   );
-  const getChatMessageMemoized = useGetChatMemoized();
+  const getChatMessageMemoized = useGetChatMessageMemoized();
 
   const isCompletedStream = useGetChatMessage(lastMessageId, (x) => x?.isCompletedStream);
   const hasReasoning = !!reasoningMessagesLength;