refactor: implement message editing and resend functionality (#5901)

* feat: implement message editing and resend functionality with block management

* fix: move start_time_millsec initialization to onChunk for accurate timing

* feat: refactor message update thunks to separate block addition and updates

* feat: enhance MessageBlockEditor with toolbar buttons and attachment functionality

* refactor: implement message editing context and integrate with message operations

* style: adjust padding in MessageBlockEditor and related components

* refactor: remove MessageEditingContext and integrate editing logic directly in Message component

* refactor: streamline message rendering logic by using conditional rendering for MessageEditor and MessageContent

* refactor: remove redundant ipcRenderer variable in useMCPServers hook

* fix: Add mock for electron's ipcRenderer to support testing

* test: Update mocks for ipcRenderer and remove redundant window.electron mock

* fix: enhance file handling in MessageEditor with new Electron API

- Added support for file dragging with visual feedback.
- Improved file type validation during drag-and-drop and clipboard pasting.
- Displayed user notifications for unsupported file types.
- Refactored file handling logic for better clarity and maintainability.

src/renderer/__tests__/setup.ts

@@ -18,3 +18,32 @@ vi.mock('electron-log/renderer', () => {
     }
   }
 })
+
+vi.stubGlobal('window', {
+  electron: {
+    ipcRenderer: {
+      on: vi.fn(), // Mocking ipcRenderer.on
+      send: vi.fn() // Mocking ipcRenderer.send
+    }
+  },
+  api: {
+    file: {
+      read: vi.fn().mockResolvedValue('[]'), // Mock file.read to return an empty array (you can customize this)
+      writeWithId: vi.fn().mockResolvedValue(undefined) // Mock file.writeWithId to do nothing
+    }
+  }
+})
+
+vi.mock('axios', () => ({
+  default: {
+    get: vi.fn().mockResolvedValue({ data: {} }), // Mocking axios GET request
+    post: vi.fn().mockResolvedValue({ data: {} }) // Mocking axios POST request
+    // You can add other axios methods like put, delete etc. as needed
+  }
+}))
+
+vi.stubGlobal('window', {
+  ...global.window, // Copy other global properties
+  addEventListener: vi.fn(), // Mock addEventListener
+  removeEventListener: vi.fn() // You can also mock removeEventListener if needed
+})

src/renderer/src/context/MessageEditingContext.tsx

@@ -0,0 +1,33 @@
+import { createContext, ReactNode, use, useState } from 'react'
+
+interface MessageEditingContextType {
+  editingMessageId: string | null
+  startEditing: (messageId: string) => void
+  stopEditing: () => void
+}
+
+const MessageEditingContext = createContext<MessageEditingContextType | null>(null)
+
+export function MessageEditingProvider({ children }: { children: ReactNode }) {
+  const [editingMessageId, setEditingMessageId] = useState<string | null>(null)
+
+  const startEditing = (messageId: string) => {
+    setEditingMessageId(messageId)
+  }
+
+  const stopEditing = () => {
+    setEditingMessageId(null)
+  }
+
+  return (
+    <MessageEditingContext value={{ editingMessageId, startEditing, stopEditing }}>{children}</MessageEditingContext>
+  )
+}
+
+export function useMessageEditing() {
+  const context = use(MessageEditingContext)
+  if (!context) {
+    throw new Error('useMessageEditing must be used within a MessageEditingProvider')
+  }
+  return context
+}

src/renderer/src/hooks/useMCPServers.ts

@@ -4,13 +4,11 @@ import { addMCPServer, deleteMCPServer, setMCPServers, updateMCPServer } from '@
 import { MCPServer } from '@renderer/types'
 import { IpcChannel } from '@shared/IpcChannel'
 
-const ipcRenderer = window.electron.ipcRenderer
-
 // Listen for server changes from main process
-ipcRenderer.on(IpcChannel.Mcp_ServersChanged, (_event, servers) => {
+window.electron.ipcRenderer.on(IpcChannel.Mcp_ServersChanged, (_event, servers) => {
   store.dispatch(setMCPServers(servers))
 })
-ipcRenderer.on(IpcChannel.Mcp_AddServer, (_event, server: MCPServer) => {
+window.electron.ipcRenderer.on(IpcChannel.Mcp_AddServer, (_event, server: MCPServer) => {
   store.dispatch(addMCPServer(server))
 })
 

src/renderer/src/hooks/useMessageOperations.ts

@@ -3,7 +3,7 @@ import Logger from '@renderer/config/logger'
 import { EVENT_NAMES, EventEmitter } from '@renderer/services/EventService'
 import { estimateUserPromptUsage } from '@renderer/services/TokenService'
 import store, { type RootState, useAppDispatch, useAppSelector } from '@renderer/store'
-import { messageBlocksSelectors, updateOneBlock } from '@renderer/store/messageBlock'
+import { updateOneBlock } from '@renderer/store/messageBlock'
 import { newMessagesActions, selectMessagesForTopic } from '@renderer/store/newMessage'
 import {
   appendAssistantResponseThunk,
@@ -13,6 +13,7 @@ import {
   deleteSingleMessageThunk,
   initiateTranslationThunk,
   regenerateAssistantResponseThunk,
+  removeBlocksThunk,
   resendMessageThunk,
   resendUserMessageWithEditThunk,
   updateMessageAndBlocksThunk,
@@ -22,21 +23,8 @@ import type { Assistant, Model, Topic } from '@renderer/types'
 import type { Message, MessageBlock } from '@renderer/types/newMessage'
 import { MessageBlockStatus, MessageBlockType } from '@renderer/types/newMessage'
 import { abortCompletion } from '@renderer/utils/abortController'
-import { findFileBlocks } from '@renderer/utils/messageUtils/find'
 import { useCallback } from 'react'
 
-const findMainTextBlockId = (message: Message): string | undefined => {
-  if (!message || !message.blocks) return undefined
-  const state = store.getState()
-  for (const blockId of message.blocks) {
-    const block = messageBlocksSelectors.selectById(state, String(blockId))
-    if (block && block.type === MessageBlockType.MAIN_TEXT) {
-      return block.id
-    }
-  }
-  return undefined
-}
-
 const selectMessagesState = (state: RootState) => state.messages
 
 export const selectNewTopicLoading = createSelector(
@@ -113,36 +101,6 @@ export function useMessageOperations(topic: Topic) {
     [dispatch, topic.id]
   )
 
-  /**
-   * 在用户消息的主文本块被编辑后重新发送该消息。 / Resends a user message after its main text block has been edited.
-   * Dispatches resendUserMessageWithEditThunk.
-   */
-  const resendUserMessageWithEdit = useCallback(
-    async (message: Message, editedContent: string, assistant: Assistant) => {
-      const mainTextBlockId = findMainTextBlockId(message)
-      if (!mainTextBlockId) {
-        console.error('Cannot resend edited message: Main text block not found.')
-        return
-      }
-
-      const files = findFileBlocks(message).map((block) => block.file)
-
-      const usage = await estimateUserPromptUsage({ content: editedContent, files })
-      const messageUpdates: Partial<Message> & Pick<Message, 'id'> = {
-        id: message.id,
-        updatedAt: new Date().toISOString(),
-        usage
-      }
-
-      await dispatch(
-        newMessagesActions.updateMessage({ topicId: topic.id, messageId: message.id, updates: messageUpdates })
-      )
-      // 对于message的修改会在下面的thunk中保存
-      await dispatch(resendUserMessageWithEditThunk(topic.id, message, mainTextBlockId, editedContent, assistant))
-    },
-    [dispatch, topic.id]
-  )
-
   /**
    * 清除当前或指定主题的所有消息。 / Clears all messages for the current or specified topic.
    * Dispatches clearTopicMessagesThunk.
@@ -309,29 +267,127 @@ export function useMessageOperations(topic: Topic) {
   )
 
   /**
-   * Updates properties of specific message blocks (e.g., content).
+   * Updates message blocks by comparing original and edited blocks.
-   * Uses the generalized thunk for persistence.
+   * Handles adding, updating, and removing blocks in a single operation.
+   * @param messageId The ID of the message to update
+   * @param editedBlocks The complete set of blocks after editing
    */
   const editMessageBlocks = useCallback(
-    async (messageId: string, updates: Partial<MessageBlock>) => {
+    async (messageId: string, editedBlocks: MessageBlock[]) => {
       if (!topic?.id) {
         console.error('[editMessageBlocks] Topic prop is not valid.')
         return
       }
 
-      const blockUpdatesListProcessed = {
+      try {
-        updatedAt: new Date().toISOString(),
+        // 1. Get the current state of the message and its blocks
-        ...updates
+        const state = store.getState()
-      }
+        const message = state.messages.entities[messageId]
+        if (!message) {
+          console.error('[editMessageBlocks] Message not found:', messageId)
+          return
+        }
 
-      const messageUpdates: Partial<Message> & Pick<Message, 'id'> = {
+        // 2. Get all original blocks
-        id: messageId,
+        const originalBlocks = message.blocks
-        updatedAt: new Date().toISOString()
+          ? (message.blocks
-      }
+              .map((blockId) => state.messageBlocks.entities[blockId])
+              .filter((block) => block !== undefined) as MessageBlock[])
+          : []
 
-      await dispatch(updateMessageAndBlocksThunk(topic.id, messageUpdates, [blockUpdatesListProcessed]))
+        // 3. Create sets for efficient comparison
+        const originalBlockIds = new Set(originalBlocks.map((block) => block.id))
+        const editedBlockIds = new Set(editedBlocks.map((block) => block.id))
+
+        // 4. Identify blocks to remove, update, and add
+        const blockIdsToRemove = originalBlocks
+          .filter((block) => !editedBlockIds.has(block.id))
+          .map((block) => block.id)
+
+        const blocksToUpdate = editedBlocks
+          .filter((block) => originalBlockIds.has(block.id))
+          .map((block) => ({
+            ...block,
+            updatedAt: new Date().toISOString()
+          }))
+
+        const blocksToAdd = editedBlocks
+          .filter((block) => !originalBlockIds.has(block.id))
+          .map((block) => ({
+            ...block,
+            updatedAt: new Date().toISOString()
+          }))
+
+        // 5. Prepare message update with new block IDs
+        const updatedBlockIds = editedBlocks.map((block) => block.id)
+        const messageUpdates: Partial<Message> & Pick<Message, 'id'> = {
+          id: messageId,
+          updatedAt: new Date().toISOString(),
+          blocks: updatedBlockIds
+        }
+
+        // 6. Log operations for debugging
+        console.log('[editMessageBlocks] Operations:', {
+          blocksToRemove: blockIdsToRemove.length,
+          blocksToUpdate: blocksToUpdate.length,
+          blocksToAdd: blocksToAdd.length
+        })
+
+        // 7. Update Redux state and database
+        // First update message and add/update blocks
+        if (blocksToAdd.length > 0) {
+          await dispatch(updateMessageAndBlocksThunk(topic.id, messageUpdates, blocksToAdd))
+        }
+
+        if (blocksToUpdate.length > 0) {
+          await dispatch(updateMessageAndBlocksThunk(topic.id, messageUpdates, blocksToUpdate))
+        }
+
+        // Then remove blocks if needed
+        if (blockIdsToRemove.length > 0) {
+          await dispatch(removeBlocksThunk(topic.id, messageId, blockIdsToRemove))
+        }
+      } catch (error) {
+        console.error('[editMessageBlocks] Failed to update message blocks:', error)
+      }
     },
-    [dispatch, topic.id]
+    [dispatch, topic?.id]
+  )
+
+  /**
+   * 在用户消息的主文本块被编辑后重新发送该消息。 / Resends a user message after its main text block has been edited.
+   * Dispatches resendUserMessageWithEditThunk.
+   */
+  const resendUserMessageWithEdit = useCallback(
+    async (message: Message, editedBlocks: MessageBlock[], assistant: Assistant) => {
+      await editMessageBlocks(message.id, editedBlocks)
+
+      const mainTextBlock = editedBlocks.find((block) => block.type === MessageBlockType.MAIN_TEXT)
+      if (!mainTextBlock) {
+        console.error('[resendUserMessageWithEdit] Main text block not found in edited blocks')
+        return
+      }
+
+      const fileBlocks = editedBlocks.filter(
+        (block) => block.type === MessageBlockType.FILE || block.type === MessageBlockType.IMAGE
+      )
+
+      const files = fileBlocks.map((block) => block.file).filter((file) => file !== undefined)
+
+      const usage = await estimateUserPromptUsage({ content: mainTextBlock.content, files })
+      const messageUpdates: Partial<Message> & Pick<Message, 'id'> = {
+        id: message.id,
+        updatedAt: new Date().toISOString(),
+        usage
+      }
+
+      await dispatch(
+        newMessagesActions.updateMessage({ topicId: topic.id, messageId: message.id, updates: messageUpdates })
+      )
+      // 对于message的修改会在下面的thunk中保存
+      await dispatch(resendUserMessageWithEditThunk(topic.id, message, assistant))
+    },
+    [dispatch, editMessageBlocks, topic.id]
   )
 
   /**

src/renderer/src/pages/home/Inputbar/AttachmentPreview.tsx

@@ -32,7 +32,55 @@ function truncateFileName(name: string, maxLength: number = MAX_FILENAME_DISPLAY
   return name.slice(0, maxLength - 3) + '...'
 }
 
-const FileNameRender: FC<{ file: FileType }> = ({ file }) => {
+export const getFileIcon = (type?: string) => {
+  if (!type) return <FileUnknownFilled />
+
+  const ext = type.toLowerCase()
+
+  if (['.jpg', '.jpeg', '.png', '.gif', '.bmp', '.webp'].includes(ext)) {
+    return <FileImageFilled />
+  }
+
+  if (['.doc', '.docx'].includes(ext)) {
+    return <FileWordFilled />
+  }
+  if (['.xls', '.xlsx'].includes(ext)) {
+    return <FileExcelFilled />
+  }
+  if (['.ppt', '.pptx'].includes(ext)) {
+    return <FilePptFilled />
+  }
+  if (ext === '.pdf') {
+    return <FilePdfFilled />
+  }
+  if (['.md', '.markdown'].includes(ext)) {
+    return <FileMarkdownFilled />
+  }
+
+  if (['.zip', '.rar', '.7z', '.tar', '.gz'].includes(ext)) {
+    return <FileZipFilled />
+  }
+
+  if (['.txt', '.json', '.log', '.yml', '.yaml', '.xml', '.csv'].includes(ext)) {
+    return <FileTextFilled />
+  }
+
+  if (['.url'].includes(ext)) {
+    return <LinkOutlined />
+  }
+
+  if (['.sitemap'].includes(ext)) {
+    return <GlobalOutlined />
+  }
+
+  if (['.folder'].includes(ext)) {
+    return <FolderOpenFilled />
+  }
+
+  return <FileUnknownFilled />
+}
+
+export const FileNameRender: FC<{ file: FileType }> = ({ file }) => {
   const [visible, setVisible] = useState<boolean>(false)
   const isImage = (ext: string) => {
     return ['.png', '.jpg', '.jpeg', '.gif', '.bmp', '.webp'].includes(ext)
@@ -85,54 +133,6 @@ const FileNameRender: FC<{ file: FileType }> = ({ file }) => {
 }
 
 const AttachmentPreview: FC<Props> = ({ files, setFiles }) => {
-  const getFileIcon = (type?: string) => {
-    if (!type) return <FileUnknownFilled />
-
-    const ext = type.toLowerCase()
-
-    if (['.jpg', '.jpeg', '.png', '.gif', '.bmp', '.webp'].includes(ext)) {
-      return <FileImageFilled />
-    }
-
-    if (['.doc', '.docx'].includes(ext)) {
-      return <FileWordFilled />
-    }
-    if (['.xls', '.xlsx'].includes(ext)) {
-      return <FileExcelFilled />
-    }
-    if (['.ppt', '.pptx'].includes(ext)) {
-      return <FilePptFilled />
-    }
-    if (ext === '.pdf') {
-      return <FilePdfFilled />
-    }
-    if (['.md', '.markdown'].includes(ext)) {
-      return <FileMarkdownFilled />
-    }
-
-    if (['.zip', '.rar', '.7z', '.tar', '.gz'].includes(ext)) {
-      return <FileZipFilled />
-    }
-
-    if (['.txt', '.json', '.log', '.yml', '.yaml', '.xml', '.csv'].includes(ext)) {
-      return <FileTextFilled />
-    }
-
-    if (['.url'].includes(ext)) {
-      return <LinkOutlined />
-    }
-
-    if (['.sitemap'].includes(ext)) {
-      return <GlobalOutlined />
-    }
-
-    if (['.folder'].includes(ext)) {
-      return <FolderOpenFilled />
-    }
-
-    return <FileUnknownFilled />
-  }
-
   if (isEmpty(files)) {
     return null
   }

src/renderer/src/pages/home/Inputbar/Inputbar.tsx

@@ -651,22 +651,6 @@ const Inputbar: FC<Props> = ({ assistant: _assistant, setActiveTopic, topic }) =
     [model, pasteLongTextAsFile, pasteLongTextThreshold, resizeTextArea, supportExts, t, text]
   )
 
-  // 添加全局粘贴事件处理
-  useEffect(() => {
-    const handleGlobalPaste = (event: ClipboardEvent) => {
-      if (document.activeElement === textareaRef.current?.resizableTextArea?.textArea) {
-        return
-      }
-
-      onPaste(event)
-    }
-
-    document.addEventListener('paste', handleGlobalPaste)
-    return () => {
-      document.removeEventListener('paste', handleGlobalPaste)
-    }
-  }, [onPaste])
-
   const handleDragOver = (e: React.DragEvent<HTMLDivElement>) => {
     e.preventDefault()
     e.stopPropagation()

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

@@ -1,12 +1,14 @@
 import ContextMenu from '@renderer/components/ContextMenu'
+import { useMessageEditing } from '@renderer/context/MessageEditingContext'
 import { useAssistant } from '@renderer/hooks/useAssistant'
+import { useMessageOperations } from '@renderer/hooks/useMessageOperations'
 import { useModel } from '@renderer/hooks/useModel'
 import { useMessageStyle, useSettings } from '@renderer/hooks/useSettings'
 import { EVENT_NAMES, EventEmitter } from '@renderer/services/EventService'
 import { getMessageModelId } from '@renderer/services/MessagesService'
 import { getModelUniqId } from '@renderer/services/ModelService'
 import { Assistant, Topic } from '@renderer/types'
-import type { Message } from '@renderer/types/newMessage'
+import type { Message, MessageBlock } from '@renderer/types/newMessage'
 import { classNames } from '@renderer/utils'
 import { Divider } from 'antd'
 import React, { Dispatch, FC, memo, SetStateAction, useCallback, useEffect, useRef } from 'react'
@@ -14,6 +16,7 @@ import { useTranslation } from 'react-i18next'
 import styled from 'styled-components'
 
 import MessageContent from './MessageContent'
+import MessageEditor from './MessageEditor'
 import MessageErrorBoundary from './MessageErrorBoundary'
 import MessageHeader from './MessageHeader'
 import MessageMenubar from './MessageMenubar'
@@ -47,11 +50,54 @@ const MessageItem: FC<Props> = ({
   const model = useModel(getMessageModelId(message), message.model?.provider) || message.model
   const { isBubbleStyle } = useMessageStyle()
   const { showMessageDivider, messageFont, fontSize } = useSettings()
+  const { editMessageBlocks, resendUserMessageWithEdit } = useMessageOperations(topic)
   const messageContainerRef = useRef<HTMLDivElement>(null)
+  const { editingMessageId, stopEditing } = useMessageEditing()
+  const isEditing = editingMessageId === message.id
+
+  useEffect(() => {
+    if (isEditing && messageContainerRef.current) {
+      messageContainerRef.current.scrollIntoView({
+        behavior: 'smooth',
+        block: 'center'
+      })
+    }
+  }, [isEditing])
+
+  const handleEditSave = useCallback(
+    async (blocks: MessageBlock[]) => {
+      try {
+        console.log('after save blocks', blocks)
+        await editMessageBlocks(message.id, blocks)
+        stopEditing()
+      } catch (error) {
+        console.error('Failed to save message blocks:', error)
+      }
+    },
+    [message, editMessageBlocks, stopEditing]
+  )
+
+  const handleEditResend = useCallback(
+    async (blocks: MessageBlock[]) => {
+      try {
+        // 编辑后重新发送消息
+        console.log('after resend blocks', blocks)
+        await resendUserMessageWithEdit(message, blocks, assistant)
+        stopEditing()
+      } catch (error) {
+        console.error('Failed to resend message:', error)
+      }
+    },
+    [message, resendUserMessageWithEdit, assistant, stopEditing]
+  )
+
+  const handleEditCancel = useCallback(() => {
+    stopEditing()
+  }, [stopEditing])
 
   const isLastMessage = index === 0
   const isAssistantMessage = message.role === 'assistant'
-  const showMenubar = !isStreaming && !message.status.includes('ing')
+  const showMenubar = !isStreaming && !message.status.includes('ing') && !isEditing
 
   const messageBorder = showMessageDivider ? undefined : 'none'
   const messageBackground = getMessageBackground(isBubbleStyle, isAssistantMessage)
@@ -114,9 +160,18 @@ const MessageItem: FC<Props> = ({
             background: messageBackground,
             overflowY: 'visible'
           }}>
-          <MessageErrorBoundary>
+          {isEditing ? (
-            <MessageContent message={message} />
+            <MessageEditor
-          </MessageErrorBoundary>
+              message={message}
+              onSave={handleEditSave}
+              onResend={handleEditResend}
+              onCancel={handleEditCancel}
+            />
+          ) : (
+            <MessageErrorBoundary>
+              <MessageContent message={message} />
+            </MessageErrorBoundary>
+          )}
           {showMenubar && (
             <MessageFooter
               className="MessageFooter"

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

@@ -20,17 +20,6 @@ const MessageContent: React.FC<Props> = ({ message }) => {
   )
 }
 
-// const SearchingContainer = styled.div`
-//   display: flex;
-//   flex-direction: row;
-//   align-items: center;
-//   background-color: var(--color-background-mute);
-//   padding: 10px;
-//   border-radius: 10px;
-//   margin-bottom: 10px;
-//   gap: 10px;
-// `
-
 const MentionTag = styled.span`
   color: var(--color-link);
 `

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

@@ -0,0 +1,367 @@
+import CustomTag from '@renderer/components/CustomTag'
+import TranslateButton from '@renderer/components/TranslateButton'
+import { isGenerateImageModel, isVisionModel } from '@renderer/config/models'
+import { useAssistant } from '@renderer/hooks/useAssistant'
+import { useSettings } from '@renderer/hooks/useSettings'
+import FileManager from '@renderer/services/FileManager'
+import { FileType, FileTypes } from '@renderer/types'
+import { Message, MessageBlock, MessageBlockStatus, MessageBlockType } from '@renderer/types/newMessage'
+import { classNames, getFileExtension } from '@renderer/utils'
+import { getFilesFromDropEvent } from '@renderer/utils/input'
+import { createFileBlock, createImageBlock } from '@renderer/utils/messageUtils/create'
+import { findAllBlocks } from '@renderer/utils/messageUtils/find'
+import { documentExts, imageExts, textExts } from '@shared/config/constant'
+import { Tooltip } from 'antd'
+import TextArea, { TextAreaRef } from 'antd/es/input/TextArea'
+import { Save, Send, X } from 'lucide-react'
+import { FC, memo, useCallback, useEffect, useMemo, useRef, useState } from 'react'
+import { useTranslation } from 'react-i18next'
+import styled from 'styled-components'
+
+import AttachmentButton, { AttachmentButtonRef } from '../Inputbar/AttachmentButton'
+import { FileNameRender, getFileIcon } from '../Inputbar/AttachmentPreview'
+import { ToolbarButton } from '../Inputbar/Inputbar'
+
+interface Props {
+  message: Message
+  onSave: (blocks: MessageBlock[]) => void
+  onResend: (blocks: MessageBlock[]) => void
+  onCancel: () => void
+}
+
+const MessageBlockEditor: FC<Props> = ({ message, onSave, onResend, onCancel }) => {
+  const allBlocks = findAllBlocks(message)
+  const [editedBlocks, setEditedBlocks] = useState<MessageBlock[]>(allBlocks)
+  const [files, setFiles] = useState<FileType[]>([])
+  const [isProcessing, setIsProcessing] = useState(false)
+  const [isFileDragging, setIsFileDragging] = useState(false)
+  const { assistant } = useAssistant(message.assistantId)
+  const model = assistant.model || assistant.defaultModel
+  const isVision = useMemo(() => isVisionModel(model), [model])
+  const supportExts = useMemo(() => [...textExts, ...documentExts, ...(isVision ? imageExts : [])], [isVision])
+  const { pasteLongTextAsFile, pasteLongTextThreshold, fontSize } = useSettings()
+  const { t } = useTranslation()
+  const textareaRef = useRef<TextAreaRef>(null)
+  const attachmentButtonRef = useRef<AttachmentButtonRef>(null)
+
+  useEffect(() => {
+    setTimeout(() => {
+      resizeTextArea()
+      if (textareaRef.current) {
+        textareaRef.current.focus({ cursor: 'end' })
+      }
+    }, 0)
+    // eslint-disable-next-line react-hooks/exhaustive-deps
+  }, [])
+
+  const resizeTextArea = useCallback(() => {
+    const textArea = textareaRef.current?.resizableTextArea?.textArea
+    if (textArea) {
+      textArea.style.height = 'auto'
+      textArea.style.height = textArea?.scrollHeight > 400 ? '400px' : `${textArea?.scrollHeight}px`
+    }
+  }, [])
+
+  const handleTextChange = (blockId: string, content: string) => {
+    setEditedBlocks((prev) => prev.map((block) => (block.id === blockId ? { ...block, content } : block)))
+  }
+
+  const onTranslated = (translatedText: string) => {
+    const mainTextBlock = editedBlocks.find((b) => b.type === MessageBlockType.MAIN_TEXT)
+    if (mainTextBlock) {
+      handleTextChange(mainTextBlock.id, translatedText)
+    }
+    setTimeout(() => resizeTextArea(), 0)
+  }
+
+  // 处理文件删除
+  const handleFileRemove = async (blockId: string) => {
+    setEditedBlocks((prev) => prev.filter((block) => block.id !== blockId))
+  }
+
+  // 处理拖拽上传
+  const handleDrop = async (e: React.DragEvent<HTMLDivElement>) => {
+    e.preventDefault()
+    e.stopPropagation()
+    setIsFileDragging(false)
+
+    const files = await getFilesFromDropEvent(e).catch((err) => {
+      console.error('[src/renderer/src/pages/home/Inputbar/Inputbar.tsx] handleDrop:', err)
+      return null
+    })
+    if (files) {
+      let supportedFiles = 0
+      files.forEach((file) => {
+        if (supportExts.includes(getFileExtension(file.path))) {
+          setFiles((prevFiles) => [...prevFiles, file])
+          supportedFiles++
+        }
+      })
+
+      // 如果有文件，但都不支持
+      if (files.length > 0 && supportedFiles === 0) {
+        window.message.info({
+          key: 'file_not_supported',
+          content: t('chat.input.file_not_supported')
+        })
+      }
+    }
+  }
+
+  const handleClick = async (withResend?: boolean) => {
+    if (isProcessing) return
+    setIsProcessing(true)
+    const updatedBlocks = [...editedBlocks]
+    if (files && files.length) {
+      const uploadedFiles = await FileManager.uploadFiles(files)
+      uploadedFiles.forEach((file) => {
+        if (file.type === FileTypes.IMAGE) {
+          const imgBlock = createImageBlock(message.id, { file, status: MessageBlockStatus.SUCCESS })
+          updatedBlocks.push(imgBlock)
+        } else {
+          const fileBlock = createFileBlock(message.id, file, { status: MessageBlockStatus.SUCCESS })
+          updatedBlocks.push(fileBlock)
+        }
+      })
+    }
+    if (withResend) {
+      onResend(updatedBlocks)
+    } else {
+      onSave(updatedBlocks)
+    }
+  }
+
+  const onPaste = useCallback(
+    async (event: ClipboardEvent) => {
+      // 1. 文本粘贴
+      const clipboardText = event.clipboardData?.getData('text')
+      if (clipboardText) {
+        if (pasteLongTextAsFile && clipboardText.length > pasteLongTextThreshold) {
+          // 长文本直接转文件，阻止默认粘贴
+          event.preventDefault()
+
+          const tempFilePath = await window.api.file.create('pasted_text.txt')
+          await window.api.file.write(tempFilePath, clipboardText)
+          const selectedFile = await window.api.file.get(tempFilePath)
+          selectedFile && setFiles((prevFiles) => [...prevFiles, selectedFile])
+          setTimeout(() => resizeTextArea(), 50)
+          return
+        }
+        // 短文本走默认粘贴行为，直接返回
+        return
+      }
+
+      // 2. 文件/图片粘贴
+      if (event.clipboardData?.files && event.clipboardData.files.length > 0) {
+        event.preventDefault()
+        for (const file of event.clipboardData.files) {
+          const filePath = window.api.file.getPathForFile(file)
+          if (!filePath) {
+            // 图像生成也支持图像编辑
+            if (file.type.startsWith('image/') && (isVisionModel(model) || isGenerateImageModel(model))) {
+              const tempFilePath = await window.api.file.create(file.name)
+              const arrayBuffer = await file.arrayBuffer()
+              const uint8Array = new Uint8Array(arrayBuffer)
+              await window.api.file.write(tempFilePath, uint8Array)
+              const selectedFile = await window.api.file.get(tempFilePath)
+              selectedFile && setFiles((prevFiles) => [...prevFiles, selectedFile])
+              break
+            } else {
+              window.message.info({
+                key: 'file_not_supported',
+                content: t('chat.input.file_not_supported')
+              })
+            }
+          }
+
+          if (supportExts.includes(getFileExtension(filePath))) {
+            const selectedFile = await window.api.file.get(filePath)
+            selectedFile && setFiles((prevFiles) => [...prevFiles, selectedFile])
+          } else {
+            window.message.info({
+              key: 'file_not_supported',
+              content: t('chat.input.file_not_supported')
+            })
+          }
+        }
+        return
+      }
+
+      // 短文本走默认粘贴行为
+    },
+    [model, pasteLongTextAsFile, pasteLongTextThreshold, resizeTextArea, supportExts, t]
+  )
+
+  const autoResizeTextArea = (e: React.ChangeEvent<HTMLTextAreaElement>) => {
+    const textarea = e.target
+    textarea.style.height = 'auto'
+    textarea.style.height = `${textarea.scrollHeight}px`
+  }
+
+  return (
+    <>
+      <EditorContainer onDragOver={(e) => e.preventDefault()} onDrop={handleDrop}>
+        {editedBlocks
+          .filter((block) => block.type === MessageBlockType.MAIN_TEXT)
+          .map((block) => (
+            <Textarea
+              className={classNames(isFileDragging && 'file-dragging')}
+              key={block.id}
+              ref={textareaRef}
+              variant="borderless"
+              value={block.content}
+              onChange={(e) => {
+                handleTextChange(block.id, e.target.value)
+                autoResizeTextArea(e)
+              }}
+              autoFocus
+              contextMenu="true"
+              spellCheck={false}
+              onPaste={(e) => onPaste(e.nativeEvent)}
+              style={{
+                fontSize,
+                padding: '0px 15px 8px 15px'
+              }}>
+              <TranslateButton onTranslated={onTranslated} />
+            </Textarea>
+          ))}
+        {(editedBlocks.some((block) => block.type === MessageBlockType.FILE || block.type === MessageBlockType.IMAGE) ||
+          files.length > 0) && (
+          <FileBlocksContainer>
+            {editedBlocks
+              .filter((block) => block.type === MessageBlockType.FILE || block.type === MessageBlockType.IMAGE)
+              .map(
+                (block) =>
+                  block.file && (
+                    <CustomTag
+                      key={block.id}
+                      icon={getFileIcon(block.file.ext)}
+                      color="#37a5aa"
+                      closable
+                      onClose={() => handleFileRemove(block.id)}>
+                      <FileNameRender file={block.file} />
+                    </CustomTag>
+                  )
+              )}
+
+            {files.map((file) => (
+              <CustomTag
+                key={file.id}
+                icon={getFileIcon(file.ext)}
+                color="#37a5aa"
+                closable
+                onClose={() => setFiles((prevFiles) => prevFiles.filter((f) => f.id !== file.id))}>
+                <FileNameRender file={file} />
+              </CustomTag>
+            ))}
+          </FileBlocksContainer>
+        )}
+
+        <ActionBar>
+          <ActionBarLeft>
+            <AttachmentButton
+              ref={attachmentButtonRef}
+              model={model}
+              files={files}
+              setFiles={setFiles}
+              ToolbarButton={ToolbarButton}
+            />
+          </ActionBarLeft>
+          <ActionBarMiddle />
+          <ActionBarRight>
+            <Tooltip title={t('common.cancel')}>
+              <ToolbarButton type="text" onClick={onCancel}>
+                <X size={16} />
+              </ToolbarButton>
+            </Tooltip>
+            <Tooltip title={t('common.save')}>
+              <ToolbarButton type="text" onClick={() => handleClick()}>
+                <Save size={16} />
+              </ToolbarButton>
+            </Tooltip>
+            <Tooltip title={t('chat.resend')}>
+              <ToolbarButton type="text" onClick={() => handleClick(true)}>
+                <Send size={16} />
+              </ToolbarButton>
+            </Tooltip>
+          </ActionBarRight>
+        </ActionBar>
+      </EditorContainer>
+    </>
+  )
+}
+
+const FileBlocksContainer = styled.div`
+  display: flex;
+  flex-wrap: wrap;
+  gap: 8px;
+  padding: 0 15px;
+  margin: 8px 0;
+  background: transplant;
+  border-radius: 4px;
+`
+
+const EditorContainer = styled.div`
+  padding: 8px 0;
+  border: 1px solid var(--color-border);
+  transition: all 0.2s ease;
+  border-radius: 15px;
+  margin-top: 0;
+  background-color: var(--color-background-opacity);
+
+  &.file-dragging {
+    border: 2px dashed #2ecc71;
+
+    &::before {
+      content: '';
+      position: absolute;
+      top: 0;
+      left: 0;
+      right: 0;
+      bottom: 0;
+      background-color: rgba(46, 204, 113, 0.03);
+      border-radius: 14px;
+      z-index: 5;
+      pointer-events: none;
+    }
+  }
+`
+
+const Textarea = styled(TextArea)`
+  padding: 0;
+  border-radius: 0;
+  display: flex;
+  flex: 1;
+  font-family: Ubuntu;
+  resize: none !important;
+  overflow: auto;
+  width: 100%;
+  box-sizing: border-box;
+  &.ant-input {
+    line-height: 1.4;
+  }
+`
+
+const ActionBar = styled.div`
+  display: flex;
+  padding: 0 8px;
+  justify-content: space-between;
+  margin-top: 8px;
+`
+
+const ActionBarLeft = styled.div`
+  display: flex;
+  align-items: center;
+`
+
+const ActionBarMiddle = styled.div`
+  flex: 1;
+`
+
+const ActionBarRight = styled.div`
+  display: flex;
+  align-items: center;
+  gap: 8px;
+`
+
+export default memo(MessageBlockEditor)

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

@@ -1,4 +1,5 @@
 import Scrollbar from '@renderer/components/Scrollbar'
+import { MessageEditingProvider } from '@renderer/context/MessageEditingContext'
 import { useMessageOperations } from '@renderer/hooks/useMessageOperations'
 import { useSettings } from '@renderer/hooks/useSettings'
 import { EVENT_NAMES, EventEmitter } from '@renderer/services/EventService'
@@ -213,34 +214,36 @@ const MessageGroup = ({ messages, topic, hidePresetMessages }: Props) => {
   )
 
   return (
-    <GroupContainer
+    <MessageEditingProvider>
-      id={`message-group-${messages[0].askId}`}
+      <GroupContainer
-      $isGrouped={isGrouped}
+        id={`message-group-${messages[0].askId}`}
-      $layout={multiModelMessageStyle}
+        $isGrouped={isGrouped}
-      className={classNames([isGrouped && 'group-container', isHorizontal && 'horizontal', isGrid && 'grid'])}>
-      <GridContainer
-        $count={messageLength}
         $layout={multiModelMessageStyle}
-        $gridColumns={gridColumns}
+        className={classNames([isGrouped && 'group-container', isHorizontal && 'horizontal', isGrid && 'grid'])}>
-        className={classNames([isGrouped && 'group-grid-container', isHorizontal && 'horizontal', isGrid && 'grid'])}>
+        <GridContainer
-        {messages.map(renderMessage)}
+          $count={messageLength}
-      </GridContainer>
+          $layout={multiModelMessageStyle}
-      {isGrouped && (
+          $gridColumns={gridColumns}
-        <MessageGroupMenuBar
+          className={classNames([isGrouped && 'group-grid-container', isHorizontal && 'horizontal', isGrid && 'grid'])}>
-          multiModelMessageStyle={multiModelMessageStyle}
+          {messages.map(renderMessage)}
-          setMultiModelMessageStyle={(style) => {
+        </GridContainer>
-            setMultiModelMessageStyle(style)
+        {isGrouped && (
-            messages.forEach((message) => {
+          <MessageGroupMenuBar
-              editMessage(message.id, { multiModelMessageStyle: style })
+            multiModelMessageStyle={multiModelMessageStyle}
-            })
+            setMultiModelMessageStyle={(style) => {
-          }}
+              setMultiModelMessageStyle(style)
-          messages={messages}
+              messages.forEach((message) => {
-          selectMessageId={selectedMessageId}
+                editMessage(message.id, { multiModelMessageStyle: style })
-          setSelectedMessage={setSelectedMessage}
+              })
-          topic={topic}
+            }}
-        />
+            messages={messages}
-      )}
+            selectMessageId={selectedMessageId}
-    </GroupContainer>
+            setSelectedMessage={setSelectedMessage}
+            topic={topic}
+          />
+        )}
+      </GroupContainer>
+    </MessageEditingProvider>
   )
 }
 

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

@@ -1,8 +1,8 @@
 import { CheckOutlined, EditOutlined, QuestionCircleOutlined, SyncOutlined } from '@ant-design/icons'
 import ObsidianExportPopup from '@renderer/components/Popups/ObsidianExportPopup'
 import SelectModelPopup from '@renderer/components/Popups/SelectModelPopup'
-import TextEditPopup from '@renderer/components/Popups/TextEditPopup'
 import { TranslateLanguageOptions } from '@renderer/config/translate'
+import { useMessageEditing } from '@renderer/context/MessageEditingContext'
 import { useMessageOperations, useTopicLoading } from '@renderer/hooks/useMessageOperations'
 import { EVENT_NAMES, EventEmitter } from '@renderer/services/EventService'
 import { getMessageTitle } from '@renderer/services/MessagesService'
@@ -23,13 +23,8 @@ import {
 } from '@renderer/utils/export'
 // import { withMessageThought } from '@renderer/utils/formats'
 import { removeTrailingDoubleSpaces } from '@renderer/utils/markdown'
-import {
+import { findMainTextBlocks, findTranslationBlocks, getMainTextContent } from '@renderer/utils/messageUtils/find'
-  findImageBlocks,
+import { Dropdown, Popconfirm, Tooltip } from 'antd'
-  findMainTextBlocks,
-  findTranslationBlocks,
-  getMainTextContent
-} from '@renderer/utils/messageUtils/find'
-import { Button, Dropdown, Popconfirm, Tooltip } from 'antd'
 import dayjs from 'dayjs'
 import { AtSign, Copy, Languages, Menu, RefreshCw, Save, Share, Split, ThumbsUp, Trash } from 'lucide-react'
 import { FilePenLine } from 'lucide-react'
@@ -65,10 +60,8 @@ const MessageMenubar: FC<Props> = (props) => {
     deleteMessage,
     resendMessage,
     regenerateAssistantMessage,
-    resendUserMessageWithEdit,
     getTranslationUpdater,
     appendAssistantResponse,
-    editMessageBlocks,
     removeMessageBlock
   } = useMessageOperations(topic)
   const loading = useTopicLoading(topic)
@@ -119,92 +112,11 @@ const MessageMenubar: FC<Props> = (props) => {
     [assistant, loading, message, resendMessage]
   )
 
+  const { startEditing } = useMessageEditing()
+
   const onEdit = useCallback(async () => {
-    // 禁用了助手消息的编辑，现在都是用户消息的编辑
+    startEditing(message.id)
-    let resendMessage = false
+  }, [message.id, startEditing])
-
-    let textToEdit = ''
-
-    const imageBlocks = findImageBlocks(message)
-    // 如果是包含图片的消息，添加图片的 markdown 格式
-    if (imageBlocks.length > 0) {
-      const imageMarkdown = imageBlocks
-        .map((image, index) => `![image-${index}](file://${image?.file?.path})`)
-        .join('\n')
-      textToEdit = `${textToEdit}\n\n${imageMarkdown}`
-    }
-    textToEdit += mainTextContent
-    // if (message.role === 'assistant' && message.model && isReasoningModel(message.model)) {
-    //   //   const processedMessage = withMessageThought(clone(message))
-    //   //   textToEdit = getMainTextContent(processedMessage)
-    //   textToEdit = mainTextContent
-    // }
-
-    const editedText = await TextEditPopup.show({
-      text: textToEdit,
-      children: (props) => {
-        const onPress = () => {
-          props.onOk?.()
-          resendMessage = true
-        }
-        return message.role === 'user' ? (
-          <ReSendButton
-            icon={<i className="iconfont icon-ic_send" style={{ color: 'var(--color-primary)' }} />}
-            onClick={onPress}>
-            {t('chat.resend')}
-          </ReSendButton>
-        ) : null
-      }
-    })
-
-    if (editedText && editedText !== textToEdit) {
-      // 解析编辑后的文本，提取图片 URL
-      // const imageRegex = /!\[image-\d+\]\((.*?)\)/g
-      // const imageUrls: string[] = []
-      // let match
-      // let content = editedText
-      // TODO 按理说图片应该走上传，不应该在这改
-      // while ((match = imageRegex.exec(editedText)) !== null) {
-      //   imageUrls.push(match[1])
-      //   content = content.replace(match[0], '')
-      // }
-      if (resendMessage) {
-        resendUserMessageWithEdit(message, editedText, assistant)
-      } else {
-        editMessageBlocks(message.id, { id: findMainTextBlocks(message)[0].id, content: editedText })
-      }
-      // // 更新消息内容，保留图片信息
-      // await editMessage(message.id, {
-      //   content: content.trim(),
-      //   metadata: {
-      //     ...message.metadata,
-      //     generateImage:
-      //       imageUrls.length > 0
-      //         ? {
-      //             type: 'url',
-      //             images: imageUrls
-      //           }
-      //         : undefined
-      //   }
-      // })
-
-      // resendMessage &&
-      //   handleResendUserMessage({
-      //     ...message,
-      //     content: content.trim(),
-      //     metadata: {
-      //       ...message.metadata,
-      //       generateImage:
-      //         imageUrls.length > 0
-      //           ? {
-      //               type: 'url',
-      //               images: imageUrls
-      //             }
-      //           : undefined
-      //     }
-      //   })
-    }
-  }, [resendUserMessageWithEdit, editMessageBlocks, assistant, mainTextContent, message, t])
 
   const handleTranslate = useCallback(
     async (language: string) => {
@@ -584,10 +496,10 @@ const ActionButton = styled.div`
   }
 `
 
-const ReSendButton = styled(Button)`
+// const ReSendButton = styled(Button)`
-  position: absolute;
+//   position: absolute;
-  top: 10px;
+//   top: 10px;
-  left: 0;
+//   left: 0;
-`
+// `
 
 export default memo(MessageMenubar)

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

@@ -971,22 +971,7 @@ export const resendMessageThunk =
  * of its associated assistant responses using resendMessageThunk.
  */
 export const resendUserMessageWithEditThunk =
-  (
+  (topicId: Topic['id'], originalMessage: Message, assistant: Assistant) => async (dispatch: AppDispatch) => {
-    topicId: Topic['id'],
-    originalMessage: Message,
-    mainTextBlockId: string,
-    editedContent: string,
-    assistant: Assistant
-  ) =>
-  async (dispatch: AppDispatch) => {
-    const blockChanges = {
-      content: editedContent,
-      updatedAt: new Date().toISOString()
-    }
-    // Update block in Redux and DB
-    dispatch(updateOneBlock({ id: mainTextBlockId, changes: blockChanges }))
-    await db.message_blocks.update(mainTextBlockId, blockChanges)
-
     // Trigger the regeneration logic for associated assistant messages
     dispatch(resendMessageThunk(topicId, originalMessage, assistant))
   }
@@ -1411,14 +1396,14 @@ 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
+    blockUpdatesList: MessageBlock[] // Block updates remain required for this thunk's purpose
   ) =>
-  async (dispatch: AppDispatch): Promise<boolean> => {
+  async (dispatch: AppDispatch): Promise<void> => {
     const messageId = messageUpdates?.id
 
     if (messageUpdates && !messageId) {
-      console.error('[updateMessageAndBlocksThunk] Message ID is required.')
+      console.error('[updateMessageAndUpdateBlocksThunk] Message ID is required.')
-      return false
+      return
     }
 
     try {
@@ -1434,14 +1419,7 @@ export const updateMessageAndBlocksThunk =
       }
 
       if (blockUpdatesList.length > 0) {
-        blockUpdatesList.forEach((blockUpdate) => {
+        dispatch(upsertManyBlocks(blockUpdatesList))
-          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)
-          }
-        })
       }
 
       // 2. 更新数据库 (在事务中)
@@ -1468,27 +1446,57 @@ export const updateMessageAndBlocksThunk =
           }
         }
 
-        // Always process block updates if the list is provided and not empty
         if (blockUpdatesList.length > 0) {
-          const validBlockUpdatesForDb = blockUpdatesList
+          await db.message_blocks.bulkPut(blockUpdatesList)
-            .map((bu) => {
+        }
-              const { id, ...changes } = bu
+      })
-              if (id && Object.keys(changes).length > 0) {
+    } catch (error) {
-                return { key: id, changes: changes }
+      console.error(`[updateMessageAndBlocksThunk] Failed to process updates for message ${messageId}:`, error)
-              }
+    }
-              return null
+  }
-            })
-            .filter((bu) => bu !== null) as { key: string; changes: Partial<MessageBlock> }[]
 
-          if (validBlockUpdatesForDb.length > 0) {
+export const removeBlocksThunk =
-            await db.message_blocks.bulkUpdate(validBlockUpdatesForDb)
+  (topicId: string, messageId: string, blockIdsToRemove: string[]) =>
-          }
+  async (dispatch: AppDispatch, getState: () => RootState): Promise<void> => {
+    if (!blockIdsToRemove.length) {
+      console.warn('[removeBlocksFromMessageThunk] No block IDs provided to remove.')
+      return
+    }
+
+    try {
+      const state = getState()
+      const message = state.messages.entities[messageId]
+
+      if (!message) {
+        console.error(`[removeBlocksFromMessageThunk] Message ${messageId} not found in state.`)
+        return
+      }
+      const blockIdsToRemoveSet = new Set(blockIdsToRemove)
+
+      const updatedBlockIds = (message.blocks || []).filter((id) => !blockIdsToRemoveSet.has(id))
+
+      // 1. Update Redux state
+      dispatch(newMessagesActions.updateMessage({ topicId, messageId, updates: { blocks: updatedBlockIds } }))
+
+      if (blockIdsToRemove.length > 0) {
+        dispatch(removeManyBlocks(blockIdsToRemove))
+      }
+
+      const finalMessagesToSave = selectMessagesForTopic(getState(), topicId)
+
+      // 2. Update database (in a transaction)
+      await db.transaction('rw', db.topics, db.message_blocks, async () => {
+        // Update the message in the topic
+        await db.topics.update(topicId, { messages: finalMessagesToSave })
+        // Delete the blocks from the database
+        if (blockIdsToRemove.length > 0) {
+          await db.message_blocks.bulkDelete(blockIdsToRemove)
         }
       })
 
-      return true
+      return
     } catch (error) {
-      console.error(`[updateMessageAndBlocksThunk] Failed to process updates for message ${messageId}:`, error)
+      console.error(`[removeBlocksFromMessageThunk] Failed to remove blocks from message ${messageId}:`, error)
-      return false
+      throw error
     }
   }

src/renderer/src/utils/messageUtils/find.ts

@@ -1,16 +1,33 @@
 import store from '@renderer/store'
 import { messageBlocksSelectors } from '@renderer/store/messageBlock'
+import { FileType } from '@renderer/types'
 import type {
   CitationMessageBlock,
   FileMessageBlock,
   ImageMessageBlock,
   MainTextMessageBlock,
   Message,
+  MessageBlock,
   ThinkingMessageBlock,
   TranslationMessageBlock
 } from '@renderer/types/newMessage'
 import { MessageBlockType } from '@renderer/types/newMessage'
 
+export const findAllBlocks = (message: Message): MessageBlock[] => {
+  if (!message || !message.blocks || message.blocks.length === 0) {
+    return []
+  }
+  const state = store.getState()
+  const allBlocks: MessageBlock[] = []
+  for (const blockId of message.blocks) {
+    const block = messageBlocksSelectors.selectById(state, blockId)
+    if (block) {
+      allBlocks.push(block)
+    }
+  }
+  return allBlocks
+}
+
 /**
  * Finds all MainTextMessageBlocks associated with a given message, in order.
  * @param message - The message object.
@@ -122,6 +139,28 @@ export const getKnowledgeBaseIds = (message: Message): string[] | undefined => {
   return firstTextBlock?.flatMap((block) => block.knowledgeBaseIds).filter((id): id is string => Boolean(id))
 }
 
+/**
+ * Gets the file content from all FileMessageBlocks and ImageMessageBlocks of a message.
+ * @param message - The message object.
+ * @returns The file content or an empty string if no file blocks are found.
+ */
+export const getFileContent = (message: Message): FileType[] => {
+  const files: FileType[] = []
+  const fileBlocks = findFileBlocks(message)
+  for (const block of fileBlocks) {
+    if (block.file) {
+      files.push(block.file)
+    }
+  }
+  const imageBlocks = findImageBlocks(message)
+  for (const block of imageBlocks) {
+    if (block.file) {
+      files.push(block.file)
+    }
+  }
+  return files
+}
+
 /**
  * Finds all CitationBlocks associated with a given message.
  * @param message - The message object.