feat(messageOperations): add editMessageBlocks functionality and update message handling logic (#5641)

feat(messageOperations): add editMessageBlocks functionality and update message handling logic

- Introduced editMessageBlocks to update properties of message blocks.
- Enhanced editMessage to include error handling and logging.
- Updated useTopicMessages to accept topic ID directly.
- Refactored MessageMenubar to utilize editMessageBlocks for editing messages.
- Improved saveMessageAndBlocksToDB for better state management.

src/renderer/src/hooks/useMessageOperations.ts

@@ -13,7 +13,8 @@ import {
   initiateTranslationThunk,
   regenerateAssistantResponseThunk,
   resendMessageThunk,
-  resendUserMessageWithEditThunk
+  resendUserMessageWithEditThunk,
+  updateMessageAndBlocksThunk
 } from '@renderer/store/thunk/messageThunk'
 import { throttledBlockDbUpdate } from '@renderer/store/thunk/messageThunk'
 import type { Assistant, Model, Topic } from '@renderer/types'
@@ -62,7 +63,7 @@ export function useMessageOperations(topic: Topic) {
     async (id: string) => {
       await dispatch(deleteSingleMessageThunk(topic.id, id))
     },
-    [dispatch, topic.id] // Use topic.id directly
+    [dispatch, topic.id]
   )
 
   /**
@@ -81,18 +82,26 @@ export function useMessageOperations(topic: Topic) {
    * 使用 newMessagesActions.updateMessage.
    */
   const editMessage = useCallback(
-    async (messageId: string, updates: Partial<Message>) => {
+    async (messageId: string, updates: Partial<Omit<Message, 'id' | 'topicId' | 'blocks'>>) => {
-      // Basic update remains the same
+      if (!topic?.id) {
-      await dispatch(newMessagesActions.updateMessage({ topicId: topic.id, messageId, updates }))
+        console.error('[editMessage] Topic prop is not valid.')
-      // TODO: Add token recalculation logic here if necessary
+        return
-      // if ('content' in updates or other relevant fields change) {
+      }
-      //   const state = store.getState(); // Need store or selector access
+      console.log(`[useMessageOperations] Editing message ${messageId} with updates:`, updates)
-      //   const message = state.messages.messagesByTopic[topic.id]?.find(m => m.id === messageId);
+
-      //   if (message) {
+      const messageUpdates: Partial<Message> & Pick<Message, 'id'> = {
-      //      const updatedUsage = await estimateTokenUsage(...); // Call estimation service
+        id: messageId,
-      //      await dispatch(newMessagesActions.updateMessage({ topicId: topic.id, messageId, updates: { usage: updatedUsage } }));
+        ...updates
-      //   }
+      }
-      // }
+
+      // Call the thunk with topic.id and only message updates
+      const success = await dispatch(updateMessageAndBlocksThunk(topic.id, messageUpdates, []))
+
+      if (success) {
+        console.log(`[useMessageOperations] Successfully edited message ${messageId} properties.`)
+      } else {
+        console.error(`[useMessageOperations] Failed to edit message ${messageId} properties.`)
+      }
     },
     [dispatch, topic.id]
   )
@@ -105,7 +114,7 @@ export function useMessageOperations(topic: Topic) {
     async (message: Message, assistant: Assistant) => {
       await dispatch(resendMessageThunk(topic.id, message, assistant))
     },
-    [dispatch, topic.id] // topic object needed by thunk
+    [dispatch, topic.id]
   )
 
   /**
@@ -122,7 +131,7 @@ export function useMessageOperations(topic: Topic) {
 
       await dispatch(resendUserMessageWithEditThunk(topic.id, message, mainTextBlockId, editedContent, assistant))
     },
-    [dispatch, topic.id] // topic object needed by thunk
+    [dispatch, topic.id]
   )
 
   /**
@@ -150,20 +159,16 @@ export function useMessageOperations(topic: Topic) {
    * 暂停当前主题正在进行的消息生成。 / Pauses ongoing message generation for the current topic.
    */
   const pauseMessages = useCallback(async () => {
-    // Use selector if preferred, but direct access is okay in callback
     const state = store.getState()
     const topicMessages = selectMessagesForTopic(state, topic.id)
     if (!topicMessages) return
 
-    // Find messages currently in progress (adjust statuses if needed)
     const streamingMessages = topicMessages.filter((m) => m.status === 'processing' || m.status === 'pending')
-
     const askIds = [...new Set(streamingMessages?.map((m) => m.askId).filter((id) => !!id) as string[])]
 
     for (const askId of askIds) {
       abortCompletion(askId)
     }
-    // Ensure loading state is set to false
     dispatch(newMessagesActions.setTopicLoading({ topicId: topic.id, loading: false }))
   }, [topic.id, dispatch])
 
@@ -172,10 +177,9 @@ export function useMessageOperations(topic: Topic) {
    */
   const resumeMessage = useCallback(
     async (message: Message, assistant: Assistant) => {
-      // Directly call the resendMessage function from this hook
       return resendMessage(message, assistant)
     },
-    [resendMessage] // Dependency is the resendMessage function itself
+    [resendMessage]
   )
 
   /**
@@ -190,7 +194,7 @@ export function useMessageOperations(topic: Topic) {
       }
       await dispatch(regenerateAssistantResponseThunk(topic.id, message, assistant))
     },
-    [dispatch, topic.id] // topic object needed by thunk
+    [dispatch, topic.id]
   )
 
   /**
@@ -209,7 +213,7 @@ export function useMessageOperations(topic: Topic) {
       }
       await dispatch(appendAssistantResponseThunk(topic.id, existingAssistantMessage.id, newModel, assistant))
     },
-    [dispatch, topic.id] // Dependencies
+    [dispatch, topic.id]
   )
 
   /**
@@ -229,7 +233,6 @@ export function useMessageOperations(topic: Topic) {
     ): Promise<((accumulatedText: string, isComplete?: boolean) => void) | null> => {
       if (!topic.id) return null
 
-      // 1. Initiate the block and get its ID
       const blockId = await dispatch(
         initiateTranslationThunk(messageId, topic.id, targetLanguage, sourceBlockId, sourceLanguage)
       )
@@ -239,23 +242,12 @@ export function useMessageOperations(topic: Topic) {
         return null
       }
 
-      // 2. Return the updater function
-      // TODO:下面这个逻辑也可以放在thunk中
       return (accumulatedText: string, isComplete: boolean = false) => {
         const status = isComplete ? MessageBlockStatus.SUCCESS : MessageBlockStatus.STREAMING
-        const changes: Partial<MessageBlock> = { content: accumulatedText, status: status } // Use Partial<MessageBlock>
+        const changes: Partial<MessageBlock> = { content: accumulatedText, status: status }
 
-        // Dispatch update to Redux store
         dispatch(updateOneBlock({ id: blockId, changes }))
-
+        throttledBlockDbUpdate(blockId, changes)
-        // Throttle update to DB
-        throttledBlockDbUpdate(blockId, changes) // Use the throttled function
-
-        // if (isComplete) {
-        //   console.log(`[TranslationUpdater] Final update for block ${blockId}.`)
-        //   // Ensure the throttled function flushes if needed, or call an immediate save
-        //   // For simplicity, we rely on the throttle's trailing call for now.
-        // }
       }
     },
     [dispatch, topic.id]
@@ -277,6 +269,38 @@ export function useMessageOperations(topic: Topic) {
     [dispatch]
   )
 
+  /**
+   * Updates properties of specific message blocks (e.g., content).
+   * Uses the generalized thunk for persistence.
+   */
+  const editMessageBlocks = useCallback(
+    // messageId?: string
+    async (blockUpdatesListRaw: Partial<MessageBlock>[]) => {
+      if (!topic?.id) {
+        console.error('[editMessageBlocks] Topic prop is not valid.')
+        return
+      }
+      if (!blockUpdatesListRaw || blockUpdatesListRaw.length === 0) {
+        console.warn('[editMessageBlocks] Received empty block updates list.')
+        return
+      }
+
+      const blockUpdatesListProcessed = blockUpdatesListRaw.map((update) => ({
+        ...update,
+        updatedAt: new Date().toISOString()
+      }))
+
+      const success = await dispatch(updateMessageAndBlocksThunk(topic.id, null, blockUpdatesListProcessed))
+
+      if (success) {
+        // console.log(`[useMessageOperations] Successfully processed block updates for message ${messageId}.`)
+      } else {
+        // console.error(`[useMessageOperations] Failed to process block updates for message ${messageId}.`)
+      }
+    },
+    [dispatch, topic.id]
+  )
+
   return {
     displayCount,
     deleteMessage,
@@ -291,12 +315,13 @@ export function useMessageOperations(topic: Topic) {
     pauseMessages,
     resumeMessage,
     getTranslationUpdater,
-    createTopicBranch
+    createTopicBranch,
+    editMessageBlocks
   }
 }
 
-export const useTopicMessages = (topic: Topic) => {
+export const useTopicMessages = (topicId: string) => {
-  const messages = useAppSelector((state) => selectMessagesForTopic(state, topic.id))
+  const messages = useAppSelector((state) => selectMessagesForTopic(state, topicId))
   return messages
 }
 

src/renderer/src/pages/home/Messages/ChatFlowHistory.tsx

@@ -616,6 +616,4 @@ const NodeContent = styled.div`
 `
 
 // 确保组件使用React.memo包装以减少不必要的重渲染
-export default memo(ChatFlowHistory, (prevProps, nextProps) => {
+export default memo(ChatFlowHistory)
-  return prevProps.conversationId === nextProps.conversationId
-})

src/renderer/src/pages/home/Messages/MessageMenubar.tsx

@@ -22,22 +22,10 @@ import {
 } from '@renderer/utils/export'
 // import { withMessageThought } from '@renderer/utils/formats'
 import { removeTrailingDoubleSpaces } from '@renderer/utils/markdown'
-import { findImageBlocks, getMainTextContent } from '@renderer/utils/messageUtils/find'
+import { findImageBlocks, findMainTextBlocks, getMainTextContent } from '@renderer/utils/messageUtils/find'
 import { Button, Dropdown, Popconfirm, Tooltip } from 'antd'
 import dayjs from 'dayjs'
-import {
+import { AtSign, Copy, Languages, Menu, RefreshCw, Save, Share, Split, ThumbsUp, Trash } from 'lucide-react'
-  AtSign,
-  Copy,
-  FilePenLine,
-  Languages,
-  Menu,
-  RefreshCw,
-  Save,
-  Share,
-  Split,
-  ThumbsUp,
-  Trash
-} from 'lucide-react'
 import { FC, memo, useCallback, useMemo, useState } from 'react'
 import { useTranslation } from 'react-i18next'
 import { useSelector } from 'react-redux'
@@ -72,7 +60,8 @@ const MessageMenubar: FC<Props> = (props) => {
     regenerateAssistantMessage,
     resendUserMessageWithEdit,
     getTranslationUpdater,
-    appendAssistantResponse
+    appendAssistantResponse,
+    editMessageBlocks
   } = useMessageOperations(topic)
   const loading = useTopicLoading(topic)
 
@@ -172,7 +161,11 @@ const MessageMenubar: FC<Props> = (props) => {
       //   imageUrls.push(match[1])
       //   content = content.replace(match[0], '')
       // }
-      resendMessage && resendUserMessageWithEdit(message, editedText, assistant)
+      if (resendMessage) {
+        resendUserMessageWithEdit(message, editedText, assistant)
+      } else {
+        editMessageBlocks([{ ...findMainTextBlocks(message)[0], content: editedText }])
+      }
       // // 更新消息内容，保留图片信息
       // await editMessage(message.id, {
       //   content: content.trim(),
@@ -204,19 +197,15 @@ const MessageMenubar: FC<Props> = (props) => {
       //     }
       //   })
     }
-  }, [resendUserMessageWithEdit, assistant, mainTextContent, message, t])
+  }, [resendUserMessageWithEdit, editMessageBlocks, assistant, mainTextContent, message, t])
 
-  // TODO 翻译
   const handleTranslate = useCallback(
     async (language: string) => {
       if (isTranslating) return
 
-      // editMessage(message.id, { translatedContent: t('translate.processing') })
-
       setIsTranslating(true)
       const messageId = message.id
       const translationUpdater = await getTranslationUpdater(messageId, language)
-      // console.log('translationUpdater', translationUpdater)
       if (!translationUpdater) return
       try {
         await translateText(mainTextContent, language, translationUpdater)
@@ -243,12 +232,12 @@ const MessageMenubar: FC<Props> = (props) => {
           window.api.file.save(fileName, mainTextContent)
         }
       },
-      {
+      // {
-        label: t('common.edit'),
+      //   label: t('common.edit'),
-        key: 'edit',
+      //   key: 'edit',
-        icon: <FilePenLine size={16} />,
+      //   icon: <FilePenLine size={16} />,
-        onClick: onEdit
+      //   onClick: onEdit
-      },
+      // },
       {
         label: t('chat.message.new.branch'),
         key: 'new-branch',
@@ -349,7 +338,7 @@ const MessageMenubar: FC<Props> = (props) => {
         ].filter(Boolean)
       }
     ],
-    [message, messageContainerRef, onEdit, onNewBranch, t, topic.name, exportMenuOptions]
+    [message, messageContainerRef, mainTextContent, onNewBranch, t, topic.name, exportMenuOptions]
   )
 
   const onRegenerate = async (e: React.MouseEvent | undefined) => {

src/renderer/src/pages/home/Messages/Messages.tsx

@@ -12,6 +12,7 @@ import { getContextCount, getGroupedMessages, getUserMessage } from '@renderer/s
 import { estimateHistoryTokens } from '@renderer/services/TokenService'
 import { useAppDispatch } from '@renderer/store'
 import { newMessagesActions } from '@renderer/store/newMessage'
+import { saveMessageAndBlocksToDB } from '@renderer/store/thunk/messageThunk'
 import type { Assistant, Topic } from '@renderer/types'
 import type { Message } from '@renderer/types/newMessage'
 import {
@@ -49,7 +50,7 @@ const Messages: React.FC<MessagesProps> = ({ assistant, topic, setActiveTopic })
   const [hasMore, setHasMore] = useState(false)
   const [isLoadingMore, setIsLoadingMore] = useState(false)
   const [isProcessingContext, setIsProcessingContext] = useState(false)
-  const messages = useTopicMessages(topic)
+  const messages = useTopicMessages(topic.id)
   const { displayCount, clearTopicMessages, deleteMessage, createTopicBranch } = useMessageOperations(topic)
   const messagesRef = useRef<Message[]>(messages)
 
@@ -147,6 +148,7 @@ const Messages: React.FC<MessagesProps> = ({ assistant, topic, setActiveTopic })
 
           const { message: clearMessage } = getUserMessage({ assistant, topic, type: 'clear' })
           dispatch(newMessagesActions.addMessage({ topicId: topic.id, message: clearMessage }))
+          await saveMessageAndBlocksToDB(clearMessage, [])
 
           scrollToBottom()
         } finally {

src/renderer/src/store/newMessage.ts

@@ -82,6 +82,7 @@ const messagesSlice = createSlice({
       const { topicId, messages } = action.payload
       messagesAdapter.upsertMany(state, messages)
       state.messageIdsByTopic[topicId] = messages.map((m) => m.id)
+      state.currentTopicId = topicId
     },
     addMessage(state, action: PayloadAction<{ topicId: string; message: Message }>) {
       const { topicId, message } = action.payload

src/renderer/src/store/thunk/messageThunk.ts

@@ -46,7 +46,7 @@ const handleChangeLoadingOfTopic = async (topicId: string) => {
   store.dispatch(newMessagesActions.setTopicLoading({ topicId, loading: false }))
 }
 
-const saveMessageAndBlocksToDB = async (message: Message, blocks: MessageBlock[]) => {
+export const saveMessageAndBlocksToDB = async (message: Message, blocks: MessageBlock[]) => {
   try {
     console.log(`[DEBUG] saveMessageAndBlocksToDB started for message ${message.id} with ${blocks.length} blocks`)
     if (blocks.length > 0) {
@@ -1397,3 +1397,112 @@ export const cloneMessagesToNewTopicThunk =
       return false // Indicate failure
     }
   }
+
+/**
+ * Thunk to edit properties of a message and/or its associated blocks.
+ * Updates Redux state and persists changes to the database within a transaction.
+ * Message updates are optional if only blocks need updating.
+ */
+export const updateMessageAndBlocksThunk =
+  (
+    topicId: string,
+    // Allow messageUpdates to be optional or just contain the ID if only blocks are updated
+    messageUpdates: (Partial<Message> & Pick<Message, 'id'>) | null, // ID is always required for context
+    blockUpdatesList: Partial<MessageBlock>[] // Block updates remain required for this thunk's purpose
+  ) =>
+  async (dispatch: AppDispatch): Promise<boolean> => {
+    const messageId = messageUpdates?.id
+    console.log(
+      `[updateMessageAndBlocksThunk] Updating message ${messageId} context in topic ${topicId}. MessageUpdates:`,
+      messageUpdates,
+      `BlockUpdates:`,
+      blockUpdatesList
+    )
+
+    if (messageUpdates && !messageId) {
+      console.error('[updateMessageAndBlocksThunk] Message ID is required.')
+      return false
+    }
+
+    try {
+      // 1. 更新 Redux Store
+      if (messageUpdates && messageId) {
+        // eslint-disable-next-line @typescript-eslint/no-unused-vars
+        const { id: msgId, ...actualMessageChanges } = messageUpdates // Separate ID from actual changes
+
+        // Only dispatch message update if there are actual changes beyond the ID
+        if (Object.keys(actualMessageChanges).length > 0) {
+          dispatch(newMessagesActions.updateMessage({ topicId, messageId, updates: actualMessageChanges }))
+          console.log(`[updateMessageAndBlocksThunk] Dispatched message property updates for ${messageId} in Redux.`)
+        } else {
+          console.log(
+            `[updateMessageAndBlocksThunk] No message property updates for ${messageId} in Redux, only processing blocks.`
+          )
+        }
+      }
+
+      if (blockUpdatesList.length > 0) {
+        blockUpdatesList.forEach((blockUpdate) => {
+          const { id: blockId, ...blockChanges } = blockUpdate
+          if (blockId && Object.keys(blockChanges).length > 0) {
+            dispatch(updateOneBlock({ id: blockId, changes: blockChanges }))
+          } else if (!blockId) {
+            console.warn('[updateMessageAndBlocksThunk] Skipping block update due to missing block ID:', blockUpdate)
+          }
+        })
+        console.log(`[updateMessageAndBlocksThunk] Dispatched ${blockUpdatesList.length} block update(s) in Redux.`)
+      }
+
+      // 2. 更新数据库 (在事务中)
+      await db.transaction('rw', db.topics, db.message_blocks, async () => {
+        // Only update topic.messages if there were actual message changes
+        if (messageUpdates && Object.keys(messageUpdates).length > 0) {
+          const topic = await db.topics.get(topicId)
+          if (topic && topic.messages) {
+            const messageIndex = topic.messages.findIndex((m) => m.id === messageId)
+            if (messageIndex !== -1) {
+              Object.assign(topic.messages[messageIndex], messageUpdates)
+              await db.topics.update(topicId, { messages: topic.messages })
+              console.log(
+                `[updateMessageAndBlocksThunk] Updated message properties for ${messageId} in DB topic ${topicId}.`
+              )
+            } else {
+              console.error(
+                `[updateMessageAndBlocksThunk] Message ${messageId} not found in DB topic ${topicId} for property update.`
+              )
+              throw new Error(`Message ${messageId} not found in DB topic ${topicId} for property update.`)
+            }
+          } else {
+            console.error(
+              `[updateMessageAndBlocksThunk] Topic ${topicId} not found or empty for message property update.`
+            )
+            throw new Error(`Topic ${topicId} not found or empty for message property update.`)
+          }
+        }
+
+        // Always process block updates if the list is provided and not empty
+        if (blockUpdatesList.length > 0) {
+          const validBlockUpdatesForDb = blockUpdatesList
+            .map((bu) => {
+              const { id, ...changes } = bu
+              if (id && Object.keys(changes).length > 0) {
+                return { key: id, changes: changes }
+              }
+              return null
+            })
+            .filter((bu) => bu !== null) as { key: string; changes: Partial<MessageBlock> }[]
+
+          if (validBlockUpdatesForDb.length > 0) {
+            await db.message_blocks.bulkUpdate(validBlockUpdatesForDb)
+            console.log(`[updateMessageAndBlocksThunk] Updated ${validBlockUpdatesForDb.length} block(s) in DB.`)
+          }
+        }
+      })
+
+      console.log(`[updateMessageAndBlocksThunk] Successfully processed updates for message ${messageId} context.`)
+      return true
+    } catch (error) {
+      console.error(`[updateMessageAndBlocksThunk] Failed to process updates for message ${messageId}:`, error)
+      return false
+    }
+  }

src/renderer/src/windows/mini/chat/components/Messages.tsx

@@ -20,7 +20,7 @@ interface ContainerProps {
 
 const Messages: FC<Props> = ({ assistant, route }) => {
   // const [messages, setMessages] = useState<Message[]>([])
-  const messages = useTopicMessages(assistant.topics[0])
+  const messages = useTopicMessages(assistant.topics[0].id)
   const containerRef = useRef<HTMLDivElement>(null)
   const messagesRef = useRef(messages)