From 95a1bdac724dcf86a7081d6eeaae6dd78f3bb9e7 Mon Sep 17 00:00:00 2001
From: kangfenmao <kangfenmao@qq.com>
Date: Fri, 24 Jan 2025 13:02:57 +0800
Subject: [PATCH] fix: resend message logic

---
 .../pages/home/Messages/MessageMenubar.tsx    | 48 ++++++++++++-------
 src/renderer/src/services/MessagesService.ts  | 17 ++++++-
 src/renderer/src/types/index.ts               |  2 +-
 3 files changed, 48 insertions(+), 19 deletions(-)

diff --git a/src/renderer/src/pages/home/Messages/MessageMenubar.tsx b/src/renderer/src/pages/home/Messages/MessageMenubar.tsx
index 8cccf936..f91979fa 100644
--- a/src/renderer/src/pages/home/Messages/MessageMenubar.tsx
+++ b/src/renderer/src/pages/home/Messages/MessageMenubar.tsx
@@ -16,11 +16,13 @@ import TextEditPopup from '@renderer/components/Popups/TextEditPopup'
 import { TranslateLanguageOptions } from '@renderer/config/translate'
 import { modelGenerating } from '@renderer/hooks/useRuntime'
 import { EVENT_NAMES, EventEmitter } from '@renderer/services/EventService'
+import { resetAssistantMessage } from '@renderer/services/MessagesService'
 import { translateText } from '@renderer/services/TranslateService'
 import { Message, Model } from '@renderer/types'
 import { removeTrailingDoubleSpaces, uuid } from '@renderer/utils'
 import { Button, Dropdown, Popconfirm, Tooltip } from 'antd'
 import dayjs from 'dayjs'
+import { isEmpty } from 'lodash'
 import { FC, useCallback, useMemo, useState } from 'react'
 import { useTranslation } from 'react-i18next'
 import styled from 'styled-components'
@@ -77,6 +79,21 @@ const MessageMenubar: FC<Props> = (props) => {
   const onResend = useCallback(async () => {
     await modelGenerating()
     const _messages = onGetMessages?.() || []
+    const groupdMessages = _messages.filter((m) => m.askId === message.id)
+
+    // Resend all groupd messages
+    if (!isEmpty(groupdMessages)) {
+      for (const assistantMessage of groupdMessages) {
+        const _model = assistantMessage.model || assistantModel
+        EventEmitter.emit(
+          EVENT_NAMES.RESEND_MESSAGE + ':' + assistantMessage.id,
+          resetAssistantMessage(assistantMessage, _model)
+        )
+      }
+      return
+    }
+
+    // If there is no groupd message, resend next message
     const index = _messages.findIndex((m) => m.id === message.id)
     const nextIndex = index + 1
     const nextMessage = _messages[nextIndex]
@@ -91,6 +108,7 @@ const MessageMenubar: FC<Props> = (props) => {
       })
     }
 
+    // If next message is not exist or next message role is user, delete current message and resend
     if (!nextMessage || nextMessage.role === 'user') {
       EventEmitter.emit(EVENT_NAMES.SEND_MESSAGE, { ...message, id: uuid() })
       onDeleteMessage?.(message)
@@ -102,20 +120,25 @@ const MessageMenubar: FC<Props> = (props) => {
 
     const editedText = await TextEditPopup.show({
       text: message.content,
-      children: (props) =>
-        message.role === 'user' ? (
+      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={() => {
-              props.onOk?.()
-              resendMessage = true
-            }}>
+            onClick={onPress}>
             {t('chat.resend')}
           </ReSendButton>
         ) : null
+      }
     })
 
-    editedText && onEditMessage?.({ ...message, content: editedText })
+    if (editedText) {
+      await onEditMessage?.({ ...message, content: editedText })
+    }
+
     resendMessage && onResend()
   }, [message, onEditMessage, onResend, t])
 
@@ -177,16 +200,7 @@ const MessageMenubar: FC<Props> = (props) => {
     const selectedModel = await SelectModelPopup.show({ model })
     if (!selectedModel) return
 
-    const _message: Message = {
-      ...message,
-      content: '',
-      reasoning_content: undefined,
-      metrics: undefined,
-      status: 'sending',
-      model: selectedModel,
-      translatedContent: undefined,
-      metadata: undefined
-    }
+    const _message: Message = resetAssistantMessage(message, selectedModel)
 
     if (message.askId && message.model) {
       return EventEmitter.emit(EVENT_NAMES.APPEND_MESSAGE, { ..._message, id: uuid() })
diff --git a/src/renderer/src/services/MessagesService.ts b/src/renderer/src/services/MessagesService.ts
index b27517a4..809c03bf 100644
--- a/src/renderer/src/services/MessagesService.ts
+++ b/src/renderer/src/services/MessagesService.ts
@@ -3,7 +3,7 @@ import { DEFAULT_CONTEXTCOUNT } from '@renderer/config/constant'
 import { getTopicById } from '@renderer/hooks/useTopic'
 import i18n from '@renderer/i18n'
 import store from '@renderer/store'
-import { Assistant, Message, Topic } from '@renderer/types'
+import { Assistant, Message, Model, Topic } from '@renderer/types'
 import { uuid } from '@renderer/utils'
 import { isEmpty, remove, takeRight } from 'lodash'
 import { NavigateFunction } from 'react-router'
@@ -153,3 +153,18 @@ export function getGroupedMessages(messages: Message[]): { [key: string]: (Messa
 export function getMessageModelId(message: Message) {
   return message?.model?.id || message.modelId
 }
+
+export function resetAssistantMessage(message: Message, model?: Model): Message {
+  return {
+    ...message,
+    model: model || message.model,
+    content: '',
+    status: 'sending',
+    translatedContent: undefined,
+    reasoning_content: undefined,
+    usage: undefined,
+    metrics: undefined,
+    metadata: undefined,
+    useful: undefined
+  }
+}
diff --git a/src/renderer/src/types/index.ts b/src/renderer/src/types/index.ts
index a382e208..c099be3c 100644
--- a/src/renderer/src/types/index.ts
+++ b/src/renderer/src/types/index.ts
@@ -53,6 +53,7 @@ export type Message = {
   createdAt: string
   status: 'sending' | 'pending' | 'success' | 'paused' | 'error'
   modelId?: string
+  model?: Model
   files?: FileType[]
   images?: string[]
   usage?: OpenAI.Completions.CompletionUsage
@@ -61,7 +62,6 @@ export type Message = {
   type: 'text' | '@' | 'clear'
   isPreset?: boolean
   mentions?: Model[]
-  model?: Model
   metadata?: {
     // Gemini
     groundingMetadata?: any