From dd5592ddbb6e8cad7d2f6082e630ca33920f5015 Mon Sep 17 00:00:00 2001
From: suyao <sy20010504@gmail.com>
Date: Sat, 20 Sep 2025 15:49:40 +0800
Subject: [PATCH] Optimize chat components with memoization and shared layout

- Wrap `SessionMessages` and `SessionInputBar` in `useMemo` to prevent unnecessary re-renders
- Refactor `AgentSessionMessages` to use shared layout components and message grouping
- Extract common styled components to `shared.tsx` for reuse across message components
---
 src/renderer/src/pages/home/Chat.tsx          | 22 ++---
 .../home/Messages/AgentSessionMessages.tsx    | 84 +++++++++----------
 .../src/pages/home/Messages/Messages.tsx      | 23 +----
 .../src/pages/home/Messages/shared.tsx        | 23 +++++
 4 files changed, 76 insertions(+), 76 deletions(-)
 create mode 100644 src/renderer/src/pages/home/Messages/shared.tsx

diff --git a/src/renderer/src/pages/home/Chat.tsx b/src/renderer/src/pages/home/Chat.tsx
index c6ab29ee3a..c3793db63c 100644
--- a/src/renderer/src/pages/home/Chat.tsx
+++ b/src/renderer/src/pages/home/Chat.tsx
@@ -17,7 +17,7 @@ import { classNames } from '@renderer/utils'
 import { Flex } from 'antd'
 import { debounce } from 'lodash'
 import { AnimatePresence, motion } from 'motion/react'
-import React, { FC, useState } from 'react'
+import React, { FC, useMemo,useState } from 'react'
 import { useHotkeys } from 'react-hotkeys-hook'
 import { useTranslation } from 'react-i18next'
 import styled from 'styled-components'
@@ -141,27 +141,27 @@ const Chat: FC<Props> = (props) => {
     ? 'calc(100vh - var(--navbar-height) - var(--navbar-height) - 12px)'
     : 'calc(100vh - var(--navbar-height))'
 
-  const SessionMessages = () => {
+  const SessionMessages = useMemo(() => {
     if (activeAgentId === null) {
-      return <div> Active Agent ID is invalid.</div>
+      return () => <div> Active Agent ID is invalid.</div>
     }
     const sessionId = activeSessionId[activeAgentId]
     if (!sessionId) {
-      return <div> Active Session ID is invalid.</div>
+      return () => <div> Active Session ID is invalid.</div>
     }
-    return <AgentSessionMessages agentId={activeAgentId} sessionId={sessionId} />
-  }
+    return () => <AgentSessionMessages agentId={activeAgentId} sessionId={sessionId} />
+  }, [activeAgentId, activeSessionId])
 
-  const SessionInputBar = () => {
+  const SessionInputBar = useMemo(() => {
     if (activeAgentId === null) {
-      return <div> Active Agent ID is invalid.</div>
+      return () => <div> Active Agent ID is invalid.</div>
     }
     const sessionId = activeSessionId[activeAgentId]
     if (!sessionId) {
-      return <div> Active Session ID is invalid.</div>
+      return () => <div> Active Session ID is invalid.</div>
     }
-    return <AgentSessionInputbar agentId={activeAgentId} sessionId={sessionId} />
-  }
+    return () => <AgentSessionInputbar agentId={activeAgentId} sessionId={sessionId} />
+  }, [activeAgentId, activeSessionId])
 
   return (
     <Container id="chat" className={classNames([messageStyle, { 'multi-select-mode': isMultiSelectMode }])}>
diff --git a/src/renderer/src/pages/home/Messages/AgentSessionMessages.tsx b/src/renderer/src/pages/home/Messages/AgentSessionMessages.tsx
index 98274160b4..9fcedf5c50 100644
--- a/src/renderer/src/pages/home/Messages/AgentSessionMessages.tsx
+++ b/src/renderer/src/pages/home/Messages/AgentSessionMessages.tsx
@@ -1,15 +1,17 @@
 import { loggerService } from '@logger'
 import ContextMenu from '@renderer/components/ContextMenu'
-import Scrollbar from '@renderer/components/Scrollbar'
 import { useSession } from '@renderer/hooks/agents/useSession'
-import Blocks from '@renderer/pages/home/Messages/Blocks'
+import { getGroupedMessages } from '@renderer/services/MessagesService'
 import { useAppSelector } from '@renderer/store'
 import { selectMessagesForTopic } from '@renderer/store/newMessage'
+import { Topic } from '@renderer/types'
 import { buildAgentSessionTopicId } from '@renderer/utils/agentSession'
-import { useMemo } from 'react'
+import { memo,useMemo } from 'react'
 import styled from 'styled-components'
 
+import MessageGroup from './MessageGroup'
 import NarrowLayout from './NarrowLayout'
+import { MessagesContainer, ScrollContainer } from './shared'
 
 const logger = loggerService.withContext('AgentSessionMessages')
 
@@ -23,6 +25,33 @@ const AgentSessionMessages: React.FC<Props> = ({ agentId, sessionId }) => {
   const sessionTopicId = useMemo(() => buildAgentSessionTopicId(sessionId), [sessionId])
   const messages = useAppSelector((state) => selectMessagesForTopic(state, sessionTopicId))
 
+  const displayMessages = useMemo(() => {
+    if (!messages || messages.length === 0) return []
+    return [...messages].reverse()
+  }, [messages])
+
+  const groupedMessages = useMemo(() => {
+    if (!displayMessages || displayMessages.length === 0) return []
+    return Object.entries(getGroupedMessages(displayMessages))
+  }, [displayMessages])
+
+  const sessionAssistantId = session?.agent_id ?? agentId
+  const sessionName = session?.name ?? sessionId
+  const sessionCreatedAt = session?.created_at ?? session?.updated_at ?? FALLBACK_TIMESTAMP
+  const sessionUpdatedAt = session?.updated_at ?? session?.created_at ?? FALLBACK_TIMESTAMP
+
+  const derivedTopic = useMemo<Topic>(
+    () => ({
+      id: sessionTopicId,
+      assistantId: sessionAssistantId,
+      name: sessionName,
+      createdAt: sessionCreatedAt,
+      updatedAt: sessionUpdatedAt,
+      messages: []
+    }),
+    [sessionTopicId, sessionAssistantId, sessionName, sessionCreatedAt, sessionUpdatedAt]
+  )
+
   logger.silly('Rendering agent session messages', {
     sessionId,
     messageCount: messages.length
@@ -33,15 +62,13 @@ const AgentSessionMessages: React.FC<Props> = ({ agentId, sessionId }) => {
       <NarrowLayout style={{ display: 'flex', flexDirection: 'column-reverse' }}>
         <ContextMenu>
           <ScrollContainer>
-            {messages
-              .slice()
-              .reverse()
-              .map((message) => (
-                <MessageRow key={message.id} $role={message.role}>
-                  <Blocks blocks={message.blocks ?? []} message={message} />
-                </MessageRow>
-              ))}
-            {!messages.length && <EmptyState>{session ? 'No messages yet.' : 'Loading session...'}</EmptyState>}
+            {groupedMessages.length > 0 ? (
+              groupedMessages.map(([key, groupMessages]) => (
+                <MessageGroup key={key} messages={groupMessages} topic={derivedTopic} />
+              ))
+            ) : (
+              <EmptyState>{session ? 'No messages yet.' : 'Loading session...'}</EmptyState>
+            )}
           </ScrollContainer>
         </ContextMenu>
       </NarrowLayout>
@@ -49,25 +76,6 @@ const AgentSessionMessages: React.FC<Props> = ({ agentId, sessionId }) => {
   )
 }
 
-const ScrollContainer = styled.div`
-  display: flex;
-  flex-direction: column-reverse;
-  gap: 12px;
-  padding: 10px 10px 20px;
-  .multi-select-mode & {
-    padding-bottom: 60px;
-  }
-`
-
-const MessageRow = styled.div<{ $role: string }>`
-  display: flex;
-  flex-direction: column;
-  align-items: ${(props) => (props.$role === 'user' ? 'flex-end' : 'flex-start')};
-  .block-wrapper {
-    max-width: 700px;
-  }
-`
-
 const EmptyState = styled.div`
   color: var(--color-text-3);
   font-size: 12px;
@@ -75,16 +83,6 @@ const EmptyState = styled.div`
   padding: 20px 0;
 `
 
-interface ContainerProps {
-  $right?: boolean
-}
+const FALLBACK_TIMESTAMP = '1970-01-01T00:00:00.000Z'
 
-const MessagesContainer = styled(Scrollbar)<ContainerProps>`
-  display: flex;
-  flex-direction: column-reverse;
-  overflow-x: hidden;
-  z-index: 1;
-  position: relative;
-`
-
-export default AgentSessionMessages
+export default memo(AgentSessionMessages)
diff --git a/src/renderer/src/pages/home/Messages/Messages.tsx b/src/renderer/src/pages/home/Messages/Messages.tsx
index 340e984a2c..7bf037411e 100644
--- a/src/renderer/src/pages/home/Messages/Messages.tsx
+++ b/src/renderer/src/pages/home/Messages/Messages.tsx
@@ -1,7 +1,6 @@
 import { loggerService } from '@logger'
 import ContextMenu from '@renderer/components/ContextMenu'
 import { LoadingIcon } from '@renderer/components/Icons'
-import Scrollbar from '@renderer/components/Scrollbar'
 import { LOAD_MORE_COUNT } from '@renderer/config/constant'
 import { useAssistant } from '@renderer/hooks/useAssistant'
 import { useChatContext } from '@renderer/hooks/useChatContext'
@@ -41,6 +40,7 @@ import MessageAnchorLine from './MessageAnchorLine'
 import MessageGroup from './MessageGroup'
 import NarrowLayout from './NarrowLayout'
 import Prompt from './Prompt'
+import { MessagesContainer, ScrollContainer } from './shared'
 
 interface MessagesProps {
   assistant: Assistant
@@ -392,25 +392,4 @@ const LoaderContainer = styled.div`
   pointer-events: none;
 `
 
-const ScrollContainer = styled.div`
-  display: flex;
-  flex-direction: column-reverse;
-  padding: 10px 10px 20px;
-  .multi-select-mode & {
-    padding-bottom: 60px;
-  }
-`
-
-interface ContainerProps {
-  $right?: boolean
-}
-
-const MessagesContainer = styled(Scrollbar)<ContainerProps>`
-  display: flex;
-  flex-direction: column-reverse;
-  overflow-x: hidden;
-  z-index: 1;
-  position: relative;
-`
-
 export default Messages
diff --git a/src/renderer/src/pages/home/Messages/shared.tsx b/src/renderer/src/pages/home/Messages/shared.tsx
new file mode 100644
index 0000000000..e9f81bbc11
--- /dev/null
+++ b/src/renderer/src/pages/home/Messages/shared.tsx
@@ -0,0 +1,23 @@
+import Scrollbar from '@renderer/components/Scrollbar'
+import styled from 'styled-components'
+
+export const ScrollContainer = styled.div`
+  display: flex;
+  flex-direction: column-reverse;
+  padding: 10px 10px 20px;
+  .multi-select-mode & {
+    padding-bottom: 60px;
+  }
+`
+
+interface ContainerProps {
+  $right?: boolean
+}
+
+export const MessagesContainer = styled(Scrollbar)<ContainerProps>`
+  display: flex;
+  flex-direction: column-reverse;
+  overflow-x: hidden;
+  z-index: 1;
+  position: relative;
+`